diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 6e1c0ac472..3448c1363d 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -68,6 +68,9 @@ type BifrostRequest struct { // Provider config must be available for each fallback's provider in account's GetConfigForProvider, // else it will be skipped. Fallbacks []Fallback `json:"fallbacks,omitempty"` + + // MCPTools contains approved MCP tools to execute (used with safe mode) + MCPTools []Tool `json:"mcp_tools,omitempty"` } // Fallback represents a fallback model to be used if the primary model is not available. @@ -401,12 +404,13 @@ type BifrostResponseChoice struct { // BifrostResponseExtraFields contains additional fields in a response. type BifrostResponseExtraFields struct { - Provider ModelProvider `json:"provider"` - Params ModelParameters `json:"model_params"` - Latency *float64 `json:"latency,omitempty"` - ChatHistory *[]BifrostMessage `json:"chat_history,omitempty"` - BilledUsage *BilledLLMUsage `json:"billed_usage,omitempty"` - RawResponse interface{} `json:"raw_response"` + Provider ModelProvider `json:"provider"` + Params ModelParameters `json:"model_params"` + Latency *float64 `json:"latency,omitempty"` + ChatHistory *[]BifrostMessage `json:"chat_history,omitempty"` + BilledUsage *BilledLLMUsage `json:"billed_usage,omitempty"` + RawResponse interface{} `json:"raw_response"` + PendingMCPTools []PendingMCPTool `json:"pending_mcp_tools,omitempty"` } const ( @@ -431,3 +435,10 @@ type ErrorField struct { Param interface{} `json:"param,omitempty"` EventID *string `json:"event_id,omitempty"` } + +// PendingMCPTool represents an MCP tool that requires user approval +type PendingMCPTool struct { + ClientName string `json:"client_name"` + Tool Tool `json:"tool"` + ToolCall ToolCall `json:"tool_call"` // Original tool call for reference +} diff --git a/plugins/mcp/README.md b/plugins/mcp/README.md new file mode 100644 index 0000000000..7f10d79b1d --- /dev/null +++ b/plugins/mcp/README.md @@ -0,0 +1,1146 @@ +# Bifrost MCP Plugin + +The **Bifrost MCP (Model Context Protocol) Plugin** provides seamless integration between Bifrost and MCP servers, enabling dynamic tool discovery, registration, and execution from both local and external MCP sources. + +## Table of Contents + +- [Overview](#overview) +- [Features](#features) +- [Quick Start](#quick-start) +- [Configuration](#configuration) +- [Usage Examples](#usage-examples) +- [Architecture](#architecture) +- [API Reference](#api-reference) +- [Advanced Features](#advanced-features) +- [Troubleshooting](#troubleshooting) +- [Contributing](#contributing) + +## Overview + +The MCP Plugin acts as a bridge between Bifrost and the Model Context Protocol ecosystem, allowing you to: + +- **Host Local Tools**: Register Go functions as MCP tools directly in Bifrost +- **Connect to External MCP Servers**: Integrate with existing MCP servers via HTTP or STDIO +- **Intelligent Tool Execution**: Support for both automatic execution and user approval workflows +- **Agentic Mode**: Enable LLM-driven tool result synthesis for natural language responses +- **Dynamic Tool Discovery**: Automatically discover and register tools from connected MCP servers + +## Features + +### πŸ”§ **Tool Management** + +- **Local Tool Hosting**: Register typed Go functions as MCP tools +- **External Tool Integration**: Connect to HTTP or STDIO-based MCP servers +- **Dynamic Discovery**: Automatically discover tools from external servers +- **Tool Filtering**: Include/exclude specific tools or clients per request + +### πŸ”’ **Security & Control** + +- **Execution Policies**: Configure per-tool approval requirements +- **Client Filtering**: Control which MCP clients are active per request +- **Tool Skipping**: Exclude specific tools from specific clients +- **Safe Defaults**: Require approval by default for external tools + +### πŸ€– **Agentic Mode** + +- **Result Synthesis**: Send tool results back to LLM for natural language responses +- **Conversation Flow**: Maintain context across tool executions +- **Fallback Handling**: Gracefully degrade when agentic mode fails + +### πŸ”Œ **Connection Types** + +- **HTTP**: Connect to web-based MCP servers +- **STDIO**: Launch and communicate with command-line MCP tools +- **Process Management**: Automatic cleanup of STDIO processes + +## Quick Start + +### 1. Basic Setup + +```go +package main + +import ( + "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/plugins/mcp" +) + +func main() { + // Create plugin with default configuration + plugin, err := mcp.NewMCPPlugin(mcp.MCPPluginConfig{ + ServerPort: ":8181", + AgenticMode: false, + }, nil) + if err != nil { + panic(err) + } + + // Create Bifrost instance + bifrost := bifrost.NewBifrost(bifrost.BifrostConfig{ + Plugins: []schemas.Plugin{plugin}, + }) + + // Register the plugin with Bifrost for agentic mode (optional) + plugin.SetBifrostClient(bifrost) +} +``` + +### 2. Register a Simple Tool + +```go +// Define tool arguments structure +type EchoArgs struct { + Message string `json:"message"` +} + +// Create tool schema +toolSchema := schemas.Tool{ + Type: "function", + Function: schemas.Function{ + Name: "echo", + Description: "Echo a message back to the user", + Parameters: schemas.FunctionParameters{ + Type: "object", + Properties: map[string]interface{}{ + "message": map[string]interface{}{ + "type": "string", + "description": "The message to echo back", + }, + }, + Required: []string{"message"}, + }, + }, +} + +// Register the tool +err := mcp.RegisterTool(plugin, "echo", "Echo a message", + func(args EchoArgs) (string, error) { + return fmt.Sprintf("Echo: %s", args.Message), nil + }, toolSchema, mcp.ToolExecutionPolicyAutoExecute) +``` + +### 3. Connect to External MCP Server + +```go +// Connect to HTTP-based MCP server +err := plugin.ConnectToExternalMCP(mcp.ExternalMCPConfig{ + Name: "weather-service", + ConnectionType: mcp.ConnectionTypeHTTP, + HTTPConnectionString: mcp.Ptr("http://localhost:3000"), +}) + +// Connect to STDIO-based MCP tool +err := plugin.ConnectToExternalMCP(mcp.ExternalMCPConfig{ + Name: "filesystem-tools", + ConnectionType: mcp.ConnectionTypeSTDIO, + StdioConfig: &mcp.StdioConfig{ + Command: "npx", + Args: []string{"@modelcontextprotocol/server-filesystem", "/tmp"}, + }, +}) +``` + +## Configuration + +### Plugin Configuration + +```go +type MCPPluginConfig struct { + ServerPort string `json:"server_port,omitempty"` // Port for local MCP server (default: ":8181") + AgenticMode bool `json:"agentic_mode,omitempty"` // Enable agentic flow + ClientConfigs []ClientExecutionConfig `json:"client_configs,omitempty"` // Per-client configurations +} +``` + +### Client Execution Configuration + +```go +type ClientExecutionConfig struct { + Name string // Client name + DefaultPolicy ToolExecutionPolicy // Default execution policy + ToolPolicies map[string]ToolExecutionPolicy // Per-tool policies + ToolsToSkip []string // Tools to exclude + ToolsToExecute []string // Tools to include (if specified, only these) +} +``` + +### Execution Policies + +- **`ToolExecutionPolicyRequireApproval`**: Tool requires user approval before execution +- **`ToolExecutionPolicyAutoExecute`**: Tool executes automatically without approval + +### Example Configuration + +```go +config := mcp.MCPPluginConfig{ + ServerPort: ":8181", + AgenticMode: true, + ClientConfigs: []mcp.ClientExecutionConfig{ + { + Name: "weather-service", + DefaultPolicy: mcp.ToolExecutionPolicyAutoExecute, + ToolPolicies: map[string]mcp.ToolExecutionPolicy{ + "get_weather": mcp.ToolExecutionPolicyAutoExecute, + "send_alert": mcp.ToolExecutionPolicyRequireApproval, + }, + ToolsToSkip: []string{"deprecated_tool"}, + }, + { + Name: "filesystem-tools", + DefaultPolicy: mcp.ToolExecutionPolicyRequireApproval, + ToolsToSkip: []string{"rm", "delete"}, + }, + }, +} +``` + +## Usage Examples + +### Example 1: File System Tools + +```go +// Connect to filesystem MCP server +err := plugin.ConnectToExternalMCP(mcp.ExternalMCPConfig{ + Name: "filesystem", + ConnectionType: mcp.ConnectionTypeSTDIO, + StdioConfig: &mcp.StdioConfig{ + Command: "npx", + Args: []string{"@modelcontextprotocol/server-filesystem", "/home/user/documents"}, + }, +}) + +// Configure with read-only permissions +plugin.clientMap["filesystem"].ExecutionConfig = mcp.ClientExecutionConfig{ + Name: "filesystem", + DefaultPolicy: mcp.ToolExecutionPolicyRequireApproval, + ToolPolicies: map[string]mcp.ToolExecutionPolicy{ + "read_file": mcp.ToolExecutionPolicyAutoExecute, + "list_files": mcp.ToolExecutionPolicyAutoExecute, + "write_file": mcp.ToolExecutionPolicyRequireApproval, + }, +} +``` + +### Example 2: Weather Service Integration + +```go +// Define weather tool arguments +type WeatherArgs struct { + Location string `json:"location"` + Units string `json:"units,omitempty"` +} + +// Register weather tool +weatherSchema := schemas.Tool{ + Type: "function", + Function: schemas.Function{ + Name: "get_weather", + Description: "Get current weather for a location", + Parameters: schemas.FunctionParameters{ + Type: "object", + Properties: map[string]interface{}{ + "location": map[string]interface{}{ + "type": "string", + "description": "City name or coordinates", + }, + "units": map[string]interface{}{ + "type": "string", + "description": "Temperature units (celsius/fahrenheit)", + "enum": []string{"celsius", "fahrenheit"}, + }, + }, + Required: []string{"location"}, + }, + }, +} + +err := mcp.RegisterTool(plugin, "get_weather", "Get current weather", + func(args WeatherArgs) (string, error) { + // Call external weather API + weather, err := getWeatherData(args.Location, args.Units) + if err != nil { + return "", err + } + return formatWeatherResponse(weather), nil + }, weatherSchema, mcp.ToolExecutionPolicyAutoExecute) +``` + +### Example 3: Client Filtering in Requests + +```go +// Create context with client filtering +ctx := context.Background() +ctx = context.WithValue(ctx, mcp.ContextKeyIncludeClients, []string{"weather-service"}) +// Only tools from weather-service will be available + +ctx = context.WithValue(ctx, mcp.ContextKeyExcludeClients, []string{"filesystem"}) +// All tools except filesystem tools will be available + +// Use in Bifrost request +request := &schemas.BifrostRequest{ + Provider: "openai", + Model: "gpt-4", + Input: schemas.RequestInput{ + ChatCompletionInput: &[]schemas.BifrostMessage{ + { + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ + ContentStr: bifrost.Ptr("What's the weather like today?"), + }, + }, + }, + }, +} + +response, err := bifrost.ChatCompletionRequestWithContext(ctx, request) +``` + +## Architecture + +### Plugin Architecture + +``` +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ Bifrost MCP Plugin β”‚ +β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ +β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ +β”‚ β”‚ Local MCP β”‚ β”‚ External MCP β”‚ β”‚ Tool β”‚ β”‚ +β”‚ β”‚ Server β”‚ β”‚ Clients β”‚ β”‚ Execution β”‚ β”‚ +β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ Engine β”‚ β”‚ +β”‚ β”‚ - Host Tools β”‚ β”‚ - HTTP Clients β”‚ β”‚ β”‚ β”‚ +β”‚ β”‚ - HTTP Server β”‚ β”‚ - STDIO Procs β”‚ β”‚ - Policies β”‚ β”‚ +β”‚ β”‚ - Tool Reg. β”‚ β”‚ - Tool Discoveryβ”‚ β”‚ - Filtering β”‚ β”‚ +β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - Agentic β”‚ β”‚ +β”‚ β”‚ Flow β”‚ β”‚ +β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ +β”‚ β”‚ Client Manager β”‚ +β”‚ β”‚ - Connection Lifecycle β”‚ +β”‚ β”‚ - Tool Mapping β”‚ +β”‚ β”‚ - Configuration Management β”‚ +β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ Bifrost Core β”‚ +β”‚ - Request Processing β”‚ +β”‚ - Response Generation β”‚ +β”‚ - Plugin Integration β”‚ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +### Tool Execution Flow + +``` +User Request + β”‚ + β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ PreHook │───▢│ LLM Process │───▢│ PostHook β”‚ +β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ +β”‚ - Add Tools β”‚ β”‚ - Generate β”‚ β”‚ - Handle Tool β”‚ +β”‚ - Filter β”‚ β”‚ Response β”‚ β”‚ Calls β”‚ +β”‚ - Handle β”‚ β”‚ - Tool Calls β”‚ β”‚ - Apply Policy β”‚ +β”‚ Approved β”‚ β”‚ β”‚ β”‚ - Execute β”‚ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ Agentic Flow β”‚ + β”‚ (Optional) β”‚ + β”‚ β”‚ + β”‚ - Send Results β”‚ + β”‚ to LLM β”‚ + β”‚ - Synthesize β”‚ + β”‚ Response β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +### Connection Types + +#### HTTP Connections + +- Direct HTTP communication with MCP servers +- Suitable for web services and remote tools +- Automatic connection management + +#### STDIO Connections + +- Launch command-line MCP tools as child processes +- Communicate via stdin/stdout +- Automatic process lifecycle management +- Process cleanup on plugin shutdown + +## API Reference + +### Core Functions + +#### `NewMCPPlugin(config MCPPluginConfig, logger schemas.Logger) (*MCPPlugin, error)` + +Creates a new MCP plugin instance with the specified configuration. + +#### `RegisterTool[T](plugin *MCPPlugin, name, description string, handler ToolHandler[T], toolSchema schemas.Tool, policy ToolExecutionPolicy) error` + +Registers a typed Go function as an MCP tool. + +#### `ConnectToExternalMCP(config ExternalMCPConfig) error` + +Connects to an external MCP server and registers its tools. + +### Plugin Interface Methods + +#### `PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.BifrostResponse, error)` + +- Adds available MCP tools to requests +- Handles approved tool execution +- Applies client filtering + +#### `PostHook(ctx *context.Context, res *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error)` + +- Processes tool calls in LLM responses +- Applies execution policies +- Handles agentic flow + +#### `Cleanup() error` + +- Terminates STDIO processes +- Disconnects MCP clients +- Cleans up resources + +## Advanced Features + +### Agentic Mode + +When enabled, tool results are sent back to the LLM for synthesis into natural language responses: + +```go +config := mcp.MCPPluginConfig{ + AgenticMode: true, + // ... other config +} + +plugin, err := mcp.NewMCPPlugin(config, logger) +plugin.SetBifrostClient(bifrostInstance) // Required for agentic mode +``` + +**Agentic Flow:** + +1. Tool executed successfully +2. Results added to conversation history +3. Synthesis prompt added +4. New LLM request made with full context +5. Natural language response generated + +### Tool and Client Filtering + +The MCP plugin provides multiple levels of filtering to control which tools are available and how they execute. + +#### Configuration-Level Tool Filtering + +Configure which tools are available from each client at startup: + +```go +config := mcp.MCPPluginConfig{ + ClientConfigs: []mcp.ClientExecutionConfig{ + { + Name: "filesystem-tools", + DefaultPolicy: mcp.ToolExecutionPolicyRequireApproval, + + // Option 1: Exclude specific tools (blacklist approach) + ToolsToSkip: []string{"rm", "delete", "format", "chmod"}, + + // Option 2: Include only specific tools (whitelist approach) + // If ToolsToExecute is specified, ONLY these tools will be available + ToolsToExecute: []string{"read_file", "list_files", "write_file"}, + + // Per-tool execution policies + ToolPolicies: map[string]mcp.ToolExecutionPolicy{ + "read_file": mcp.ToolExecutionPolicyAutoExecute, + "write_file": mcp.ToolExecutionPolicyRequireApproval, + }, + }, + }, +} +``` + +**Configuration-Level Priority Rules:** + +1. **`ToolsToExecute` takes precedence**: If specified, only these tools are available (whitelist) +2. **`ToolsToSkip` is secondary**: Only applies when `ToolsToExecute` is empty (blacklist) +3. **Empty configurations**: All discovered tools are available + +#### Request-Level Client Filtering + +Control which MCP clients are active per individual request: + +```go +// Whitelist mode - only include specific clients +ctx = context.WithValue(ctx, mcp.ContextKeyIncludeClients, []string{"weather", "calendar"}) + +// Blacklist mode - exclude specific clients +ctx = context.WithValue(ctx, mcp.ContextKeyExcludeClients, []string{"filesystem", "admin-tools"}) + +// Use in request +response, err := bifrost.ChatCompletionRequestWithContext(ctx, request) +``` + +**Request-Level Priority Rules:** + +1. **Include takes absolute precedence**: If `ContextKeyIncludeClients` is set, only those clients are used +2. **Exclude is secondary**: Only applies when include list is empty +3. **Empty filters**: All configured clients are available + +#### Combined Example: Multi-Level Filtering + +```go +// 1. Configuration Level: Set up clients with tool filtering +config := mcp.MCPPluginConfig{ + ClientConfigs: []mcp.ClientExecutionConfig{ + { + Name: "filesystem", + DefaultPolicy: mcp.ToolExecutionPolicyRequireApproval, + ToolsToExecute: []string{"read_file", "list_files"}, // Only safe read operations + }, + { + Name: "weather", + DefaultPolicy: mcp.ToolExecutionPolicyAutoExecute, + // All tools available (no filtering) + }, + { + Name: "admin-tools", + DefaultPolicy: mcp.ToolExecutionPolicyRequireApproval, + ToolsToSkip: []string{"delete_user", "reset_system"}, // Exclude dangerous operations + }, + }, +} + +// 2. Request Level: Further filter clients per request +ctx := context.Background() + +// For safe operations - include filesystem and weather only +ctx = context.WithValue(ctx, mcp.ContextKeyIncludeClients, []string{"filesystem", "weather"}) + +// For admin operations - exclude only high-risk client +// ctx = context.WithValue(ctx, mcp.ContextKeyExcludeClients, []string{"admin-tools"}) +``` + +#### Filtering Priority Summary + +**Overall Priority Order (highest to lowest):** + +1. **Request-level include** (`ContextKeyIncludeClients`) - Absolute whitelist +2. **Request-level exclude** (`ContextKeyExcludeClients`) - Applied if no include list +3. **Config-level tool whitelist** (`ToolsToExecute`) - Per-client tool whitelist +4. **Config-level tool blacklist** (`ToolsToSkip`) - Per-client tool blacklist +5. **Default**: All tools from all clients available + +**Conflict Resolution Examples:** + +```go +// Example 1: Include overrides exclude +ctx = context.WithValue(ctx, mcp.ContextKeyIncludeClients, []string{"weather"}) +ctx = context.WithValue(ctx, mcp.ContextKeyExcludeClients, []string{"filesystem"}) +// Result: Only "weather" client tools available (exclude ignored) + +// Example 2: Config whitelist overrides blacklist +clientConfig := mcp.ClientExecutionConfig{ + Name: "filesystem", + ToolsToExecute: []string{"read_file"}, // Whitelist + ToolsToSkip: []string{"write_file"}, // Blacklist (ignored) +} +// Result: Only "read_file" available from filesystem client + +// Example 3: Layered filtering +// Config: filesystem client has only ["read_file", "write_file"] +// Request: ctx includes only ["filesystem", "weather"] +// Result: filesystem tools (read_file, write_file) + all weather tools +``` + +#### Dynamic Filtering Use Cases + +**Security Contexts:** + +```go +// High-security context - read-only operations +ctx = context.WithValue(ctx, mcp.ContextKeyIncludeClients, []string{"filesystem-readonly", "weather"}) + +// Admin context - full access except dangerous operations +ctx = context.WithValue(ctx, mcp.ContextKeyExcludeClients, []string{"destructive-tools"}) +``` + +**User Role-Based Filtering:** + +```go +func getContextForUserRole(role string) context.Context { + ctx := context.Background() + + switch role { + case "admin": + // Admins get all tools + return ctx + case "user": + // Users get safe tools only + return context.WithValue(ctx, mcp.ContextKeyIncludeClients, + []string{"weather", "calendar", "filesystem-readonly"}) + case "guest": + // Guests get minimal access + return context.WithValue(ctx, mcp.ContextKeyIncludeClients, + []string{"weather"}) + default: + // No tools for unknown roles + return context.WithValue(ctx, mcp.ContextKeyIncludeClients, []string{}) + } +} +``` + +### Tool Execution Policies + +Configure fine-grained control over tool execution: + +```go +clientConfig := mcp.ClientExecutionConfig{ + Name: "external-tools", + DefaultPolicy: mcp.ToolExecutionPolicyRequireApproval, + ToolPolicies: map[string]mcp.ToolExecutionPolicy{ + "safe_read_operation": mcp.ToolExecutionPolicyAutoExecute, + "write_operation": mcp.ToolExecutionPolicyRequireApproval, + "admin_operation": mcp.ToolExecutionPolicyRequireApproval, + }, + ToolsToSkip: []string{"deprecated_tool", "dangerous_operation"}, +} +``` + +### Safe Execution Flows + +The MCP plugin provides robust security through safe execution flows, allowing users to review and approve tool executions before they occur. This is especially important when dealing with external MCP tools that might perform sensitive operations. + +#### Understanding Tool Execution Modes + +**Auto-Execute Mode:** + +- Tools with `ToolExecutionPolicyAutoExecute` run immediately +- Best for safe, read-only operations (weather, search, etc.) +- Results are processed directly or sent to LLM for synthesis + +**Approval-Required Mode:** + +- Tools with `ToolExecutionPolicyRequireApproval` require user confirmation +- Default for external tools and potentially dangerous operations +- Users see tool details before execution + +#### Implementing Safe Execution in Your Application + +Here's how to implement the user approval workflow in your application, based on the pattern used in the chatbot example: + +```go +func handleChatMessage(session *ChatSession, message string) (string, error) { + // Send request to Bifrost + response, err := session.client.ChatCompletionRequest(context.Background(), request) + if err != nil { + return "", err + } + + // Check if response contains pending tools requiring approval + if response.ExtraFields.PendingMCPTools != nil && len(*response.ExtraFields.PendingMCPTools) > 0 { + return handlePendingTools(session, response) + } + + // Normal response processing + return extractResponseText(response), nil +} + +func handlePendingTools(session *ChatSession, response *schemas.BifrostResponse) (string, error) { + pendingTools := *response.ExtraFields.PendingMCPTools + + // Display tools to user for approval + fmt.Println("πŸ”’ Safe Mode: The following tools require your approval:") + fmt.Println("=====================================================") + + for i, tool := range pendingTools { + fmt.Printf("[%d] Tool: %s\n", i+1, tool.Tool.Function.Name) + fmt.Printf(" Client: %s\n", tool.ClientName) + fmt.Printf(" Description: %s\n", tool.Tool.Function.Description) + fmt.Printf(" Arguments: %s\n", tool.ToolCall.Function.Arguments) + fmt.Println() + } + + // Get user approval + fmt.Print("Do you want to approve these tools? (y/n): ") + if !getUserApproval() { + return "❌ Tool execution cancelled by user.", nil + } + + // Create approved tools list + approvedTools := make([]schemas.Tool, 0) + for _, pendingTool := range pendingTools { + approvedTools = append(approvedTools, pendingTool.Tool) + } + + // Reconstruct conversation for approved execution + conversationHistory := session.history + assistantMessage := response.Choices[0].Message + conversationHistory = append(conversationHistory, assistantMessage) + + // Create new request with approved tools + // No tool placeholder messages needed - the MCP plugin handles execution internally + approvedRequest := &schemas.BifrostRequest{ + Provider: session.config.Provider, + Model: session.config.Model, + Input: schemas.RequestInput{ + ChatCompletionInput: &conversationHistory, + }, + MCPTools: &approvedTools, // This signals approved execution to MCP plugin + } + + // Execute approved tools + fmt.Println("βœ… Tools approved - executing...") + finalResponse, err := session.client.ChatCompletionRequest(context.Background(), approvedRequest) + if err != nil { + return "", fmt.Errorf("approved tool execution failed: %s", err.Error.Message) + } + + // Update conversation history with final result + if len(finalResponse.Choices) > 0 { + finalMessage := finalResponse.Choices[0].Message + session.history = append(session.history, finalMessage) + return extractResponseText(finalResponse), nil + } + + return "", fmt.Errorf("no response received from approved execution") +} + +func getUserApproval() bool { + var input string + fmt.Scanln(&input) + return strings.ToLower(strings.TrimSpace(input)) == "y" || + strings.ToLower(strings.TrimSpace(input)) == "yes" +} + +func extractResponseText(response *schemas.BifrostResponse) string { + if len(response.Choices) == 0 { + return "" + } + + message := response.Choices[0].Message + if message.Content.ContentStr != nil { + return *message.Content.ContentStr + } + + // Handle content blocks + if message.Content.ContentBlocks != nil { + var textParts []string + for _, block := range *message.Content.ContentBlocks { + if block.Text != nil { + textParts = append(textParts, *block.Text) + } + } + return strings.Join(textParts, "\n") + } + + return "" +} +``` + +#### Safe Execution Flow Diagram + +The following diagram illustrates the complete safe execution flow: + +```mermaid +sequenceDiagram + participant User + participant App as "Your Application" + participant Bifrost + participant MCP as "MCP Plugin" + participant LLM as "Language Model" + participant MCPServer as "MCP Server" + + User->>App: Send message + App->>Bifrost: ChatCompletionRequest() + Bifrost->>MCP: PreHook() + MCP-->>Bifrost: Add available tools to request + Bifrost->>LLM: Request with tools + LLM-->>Bifrost: Response with tool calls + Bifrost->>MCP: PostHook() + + alt Tools require approval + MCP-->>Bifrost: Return PendingMCPTools + Bifrost-->>App: Response with pending tools + App->>User: Display tools for approval
"πŸ”’ Safe Mode: Tools require approval" + User->>App: Approve tools (y/n) + + alt User approves + App->>Bifrost: New request with MCPTools approved + Bifrost->>MCP: PreHook() with approved tools + MCP->>MCPServer: Execute approved tools + MCPServer-->>MCP: Tool results + + alt Agentic mode enabled + MCP->>Bifrost: Send conversation + results to LLM + Bifrost->>LLM: Request for synthesis + LLM-->>Bifrost: Synthesized response + Bifrost-->>App: Natural language response + else Non-agentic mode + MCP-->>Bifrost: Direct tool results + Bifrost-->>App: Tool execution results + end + + App->>User: Show final response + else User denies + App->>User: "❌ Tool execution cancelled" + end + + else Tools auto-execute + MCP->>MCPServer: Execute tools immediately + MCPServer-->>MCP: Tool results + + alt Agentic mode enabled + MCP->>Bifrost: Send results to LLM for synthesis + Bifrost->>LLM: Request for synthesis + LLM-->>Bifrost: Synthesized response + Bifrost-->>App: Natural language response + else Non-agentic mode + MCP-->>Bifrost: Direct tool results + Bifrost-->>App: Tool execution results + end + + App->>User: Show response + end +``` + +#### Key Components of Safe Execution + +**1. Pending Tools Detection** + +When the MCP plugin determines that tools require approval, it returns them in `response.ExtraFields.PendingMCPTools`: + +```go +type PendingMCPTool struct { + ClientName string // Which MCP client owns this tool + Tool schemas.Tool // Tool definition and schema + ToolCall schemas.ToolCall // Actual tool call with arguments +} +``` + +**2. User Approval Interface** + +Your application should display: + +- Tool name and description +- MCP client name (for trust assessment) +- Tool arguments (for security review) +- Clear approve/deny options + +**3. Approved Execution Request** + +When tools are approved, create a new request with: + +- Original conversation history +- Assistant message with tool calls +- `MCPTools` field containing approved tools (no placeholder messages needed) + +**4. Conversation Flow Management** + +The plugin handles different conversation flows: + +**Agentic Mode (Recommended):** + +- Tool results are sent back to the LLM +- LLM synthesizes results into natural language +- Maintains conversational context +- Provides better user experience + +**Non-Agentic Mode:** + +- Tool results returned directly +- Raw tool output shown to user +- Less context-aware responses + +#### Security Best Practices + +**1. Default to Approval-Required** + +```go +// Secure by default - external tools require approval +clientConfig := mcp.ClientExecutionConfig{ + Name: "external-tools", + DefaultPolicy: mcp.ToolExecutionPolicyRequireApproval, + // Only explicitly mark safe tools for auto-execution + ToolPolicies: map[string]mcp.ToolExecutionPolicy{ + "get_weather": mcp.ToolExecutionPolicyAutoExecute, + "search_web": mcp.ToolExecutionPolicyAutoExecute, + "file_operation": mcp.ToolExecutionPolicyRequireApproval, + }, +} +``` + +**2. Clear Tool Information Display** + +Always show users: + +- What the tool does (description) +- What data it will access (arguments) +- Which system it comes from (client name) +- Potential risks or side effects + +**3. Audit Trail** + +Consider logging tool approvals and executions: + +```go +func logToolExecution(tool schemas.Tool, approved bool, userID string) { + log.Printf("Tool execution: %s, Approved: %t, User: %s, Args: %v", + tool.Function.Name, approved, userID, tool.Function.Parameters) +} +``` + +**4. Granular Control** + +Use filtering to control tool availability by context: + +```go +// Restrict tools based on user role or request context +func getSecureContext(userRole string) context.Context { + ctx := context.Background() + + switch userRole { + case "admin": + // Admins can access all tools + return ctx + case "user": + // Regular users get safe tools only + return context.WithValue(ctx, mcp.ContextKeyIncludeClients, + []string{"weather", "search", "calendar"}) + case "readonly": + // Read-only users get minimal access + return context.WithValue(ctx, mcp.ContextKeyIncludeClients, + []string{"weather", "search"}) + default: + // Unknown roles get no tools + return context.WithValue(ctx, mcp.ContextKeyIncludeClients, []string{}) + } +} +``` + +#### Example: Complete Safe Execution Implementation + +Here's a complete example showing safe execution in a web application: + +```go +func (h *ChatHandler) handleToolApproval(w http.ResponseWriter, r *http.Request) { + var request struct { + PendingTools []schemas.PendingMCPTool `json:"pending_tools"` + Approved bool `json:"approved"` + SessionID string `json:"session_id"` + } + + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + http.Error(w, "Invalid request", http.StatusBadRequest) + return + } + + session := h.getSession(request.SessionID) + if session == nil { + http.Error(w, "Session not found", http.StatusNotFound) + return + } + + if !request.Approved { + // User denied tool execution + response := map[string]interface{}{ + "message": "❌ Tool execution cancelled by user", + "type": "cancellation", + } + json.NewEncoder(w).Encode(response) + return + } + + // Log the approval for audit + h.logToolApproval(request.PendingTools, session.UserID) + + // Execute approved tools + result, err := h.executeApprovedTools(session, request.PendingTools) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + response := map[string]interface{}{ + "message": result, + "type": "execution_result", + } + json.NewEncoder(w).Encode(response) +} +``` + +This safe execution pattern ensures that users maintain control over tool execution while providing a smooth experience for approved operations. + +## Troubleshooting + +### Common Issues + +#### 1. Connection Failures + +**STDIO Connection Issues:** + +``` +Error: failed to start command 'npx @modelcontextprotocol/server-filesystem' +``` + +**Solutions:** + +- Verify the command exists and is executable +- Check command arguments are correct +- Ensure required dependencies are installed +- Check file permissions + +**HTTP Connection Issues:** + +``` +Error: failed to initialize external MCP client: connection refused +``` + +**Solutions:** + +- Verify the HTTP server is running +- Check the URL is correct and accessible +- Verify network connectivity +- Check firewall settings + +#### 2. Tool Registration Failures + +**Tool Already Exists:** + +``` +Error: tool 'echo' already registered +``` + +**Solutions:** + +- Use unique tool names +- Check for duplicate registrations +- Clear existing tools if needed + +#### 3. Agentic Mode Issues + +**Bifrost Client Not Set:** + +``` +Warning: Agentic mode is enabled but Bifrost client is not set +``` + +**Solutions:** + +```go +plugin.SetBifrostClient(bifrostInstance) +``` + +#### 4. Tool Filtering Issues + +**No Tools Available:** + +``` +Error: No MCP tools found in response +``` + +**Common Causes & Solutions:** + +- **Over-restrictive filtering**: Check if `ContextKeyIncludeClients` is too narrow +- **All tools skipped**: Review `ToolsToSkip` configuration for each client +- **Client connection issues**: Verify external MCP clients are connected +- **Empty whitelist**: If `ToolsToExecute` is set but empty, no tools will be available + +```go +// Debug tool availability +tools := plugin.getFilteredAvailableTools(&ctx) +fmt.Printf("Available tools: %d\n", len(tools)) +for _, tool := range tools { + fmt.Printf("- %s (from client: %s)\n", tool.Function.Name, "unknown") +} +``` + +**Unexpected Tool Availability:** + +``` +Warning: Restricted tool 'delete_all_files' is available when it shouldn't be +``` + +**Solutions:** + +- **Check priority order**: Ensure `ToolsToExecute` whitelist is properly configured +- **Verify client filtering**: Make sure dangerous clients are excluded at request level +- **Review configuration**: Confirm `ToolsToSkip` is correctly specified + +**Tools Not Executing:** + +``` +Error: Tool 'safe_operation' requires approval but should auto-execute +``` + +**Solutions:** + +- **Check execution policies**: Verify `ToolPolicies` configuration +- **Review default policy**: Ensure `DefaultPolicy` is set correctly +- **Policy hierarchy**: Tool-specific policies override default policies + +```go +// Debug execution policy +policy := plugin.getToolExecutionPolicy("safe_operation", "filesystem-client") +fmt.Printf("Execution policy for safe_operation: %s\n", policy) +``` + +### Debugging Tips + +#### Enable Debug Logging + +```go +logger := bifrost.NewDefaultLogger(schemas.LogLevelDebug) +plugin, err := mcp.NewMCPPlugin(config, logger) +``` + +#### Check Tool Registration + +```go +// List available tools +tools := plugin.getFilteredAvailableTools(&ctx) +for _, tool := range tools { + fmt.Printf("Tool: %s - %s\n", tool.Function.Name, tool.Function.Description) +} +``` + +#### Debug Filtering Configuration + +```go +// Check what clients are active +ctx := context.Background() +ctx = context.WithValue(ctx, mcp.ContextKeyIncludeClients, []string{"filesystem"}) + +// Verify filtering is working +tools := plugin.getFilteredAvailableTools(&ctx) +fmt.Printf("Filtered tools count: %d\n", len(tools)) + +// Check individual client configuration +for clientName, client := range plugin.clientMap { + fmt.Printf("Client: %s\n", clientName) + fmt.Printf(" Default Policy: %s\n", client.ExecutionConfig.DefaultPolicy) + fmt.Printf(" Tools to Skip: %v\n", client.ExecutionConfig.ToolsToSkip) + fmt.Printf(" Tools to Execute: %v\n", client.ExecutionConfig.ToolsToExecute) + fmt.Printf(" Available Tools: %d\n", len(client.ToolMap)) +} +``` + +#### Monitor Process Status + +```go +// Check STDIO process status +for name, client := range plugin.clientMap { + if client.StdioCommand != nil { + fmt.Printf("Client %s: PID %d, State: %s\n", + name, client.StdioCommand.Process.Pid, client.StdioCommand.ProcessState) + } +} +``` + +--- + +For more information, see the [main Bifrost documentation](../../README.md). diff --git a/plugins/mcp/chat/chatbot.go b/plugins/mcp/chat/chatbot.go new file mode 100644 index 0000000000..2257f0b736 --- /dev/null +++ b/plugins/mcp/chat/chatbot.go @@ -0,0 +1,604 @@ +package main + +import ( + "bufio" + "context" + "fmt" + "os" + "os/signal" + "strings" + "sync" + "syscall" + "time" + + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + mcp "github.com/maximhq/bifrost/plugins/mcp" +) + +// ChatbotConfig holds configuration for the chatbot +type ChatbotConfig struct { + Provider schemas.ModelProvider + Model string + MCPAgenticMode bool + MCPServerPort string + EnableMaximMCP bool + Temperature *float64 + MaxTokens *int +} + +// ChatSession manages the conversation state +type ChatSession struct { + history []schemas.BifrostMessage + client *bifrost.Bifrost + mcpPlugin *mcp.MCPPlugin + config ChatbotConfig + systemPrompt string +} + +// BaseAccount implements the schemas.Account interface for testing purposes. +type BaseAccount struct{} + +func (baseAccount *BaseAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) { + return []schemas.ModelProvider{schemas.OpenAI}, nil +} + +func (baseAccount *BaseAccount) GetKeysForProvider(providerKey schemas.ModelProvider) ([]schemas.Key, error) { + return []schemas.Key{ + { + Value: os.Getenv("OPENAI_API_KEY"), + Models: []string{"gpt-4o-mini", "gpt-4-turbo"}, + Weight: 1.0, + }, + }, nil +} + +func (baseAccount *BaseAccount) GetConfigForProvider(providerKey schemas.ModelProvider) (*schemas.ProviderConfig, error) { + return &schemas.ProviderConfig{ + NetworkConfig: schemas.DefaultNetworkConfig, + ConcurrencyAndBufferSize: schemas.DefaultConcurrencyAndBufferSize, + }, nil +} + +// NewChatSession creates a new chat session with the given configuration +func NewChatSession(config ChatbotConfig) (*ChatSession, error) { + // Create MCP plugin with client-level execution policies + mcpConfig := mcp.MCPPluginConfig{ + ServerPort: config.MCPServerPort, + AgenticMode: config.MCPAgenticMode, + ClientConfigs: []mcp.ClientExecutionConfig{ + // All clients require approval by default for safety + { + Name: "maxim-mcp", + DefaultPolicy: mcp.ToolExecutionPolicyRequireApproval, + ToolPolicies: map[string]mcp.ToolExecutionPolicy{}, // Can override specific tools here + ToolsToSkip: []string{"get-current-utc-time", "get-maxim-workspace-id"}, // Skip these tools for this client + }, + { + Name: "serper-web-search-mcp", + DefaultPolicy: mcp.ToolExecutionPolicyAutoExecute, + ToolPolicies: map[string]mcp.ToolExecutionPolicy{}, // Can override specific tools here + ToolsToSkip: []string{}, // No tools to skip for this client + }, + { + Name: "context7", + DefaultPolicy: mcp.ToolExecutionPolicyAutoExecute, + ToolPolicies: map[string]mcp.ToolExecutionPolicy{}, // Can override specific tools here + ToolsToSkip: []string{}, // No tools to skip for this client + }, + }, + } + + mcpPlugin, err := mcp.NewMCPPlugin(mcpConfig, nil) // nil logger will use default + if err != nil { + return nil, fmt.Errorf("failed to create MCP plugin: %w", err) + } + + // Connect to external MCP servers based on config + if config.EnableMaximMCP { + fmt.Println("πŸ”Œ Connecting to Maxim MCP server...") + err = mcpPlugin.ConnectToExternalMCP(mcp.ExternalMCPConfig{ + Name: "maxim-mcp", + ConnectionType: mcp.ConnectionTypeSTDIO, + StdioConfig: &mcp.StdioConfig{ + Command: "npx", + Args: []string{"-y", "@maximai/mcp-server@latest"}, + }, + }) + if err != nil { + fmt.Printf("⚠️ Warning: Failed to connect to Maxim MCP: %v\n", err) + } else { + fmt.Println("βœ… Connected to Maxim MCP server") + } + } + + fmt.Println("πŸ”Œ Connecting to Serper MCP server...") + err = mcpPlugin.ConnectToExternalMCP(mcp.ExternalMCPConfig{ + Name: "serper-web-search-mcp", + ConnectionType: mcp.ConnectionTypeSTDIO, + StdioConfig: &mcp.StdioConfig{ + Command: "npx", + Args: []string{"-y", "serper-search-scrape-mcp-server"}, + }, + }) + if err != nil { + fmt.Printf("⚠️ Warning: Failed to connect to Serper MCP: %v\n", err) + } else { + fmt.Println("βœ… Connected to Serper MCP server") + } + + fmt.Println("πŸ”Œ Connecting to Context7 MCP server...") + err = mcpPlugin.ConnectToExternalMCP(mcp.ExternalMCPConfig{ + Name: "context7", + ConnectionType: mcp.ConnectionTypeSTDIO, + StdioConfig: &mcp.StdioConfig{ + Command: "npx", + Args: []string{"-y", "@upstash/context7-mcp"}, + }, + }) + if err != nil { + fmt.Printf("⚠️ Warning: Failed to connect to Context7 MCP: %v\n", err) + } else { + fmt.Println("βœ… Connected to Context7 MCP server") + } + + // Initialize Bifrost + account := &BaseAccount{} + client, err := bifrost.Init(schemas.BifrostConfig{ + Account: account, + Plugins: []schemas.Plugin{mcpPlugin}, + Logger: bifrost.NewDefaultLogger(schemas.LogLevelInfo), + }) + if err != nil { + return nil, fmt.Errorf("failed to initialize Bifrost: %w", err) + } + + // Set Bifrost client for agentic mode + if config.MCPAgenticMode { + mcpPlugin.SetBifrostClient(client) + } + + session := &ChatSession{ + history: make([]schemas.BifrostMessage, 0), + client: client, + mcpPlugin: mcpPlugin, + config: config, + systemPrompt: "You are a helpful AI assistant with access to various tools. " + + "Use the available tools when they can help answer the user's questions more accurately or provide additional information.", + } + + // Add system message to history + if session.systemPrompt != "" { + session.history = append(session.history, schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleSystem, + Content: schemas.MessageContent{ + ContentStr: &session.systemPrompt, + }, + }) + } + + return session, nil +} + +// AddUserMessage adds a user message to the conversation history +func (s *ChatSession) AddUserMessage(message string) { + userMessage := schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ + ContentStr: &message, + }, + } + s.history = append(s.history, userMessage) +} + +// SendMessage sends a message and returns the assistant's response +func (s *ChatSession) SendMessage(message string) (string, error) { + // Add user message to history + s.AddUserMessage(message) + + // Prepare model parameters + params := &schemas.ModelParameters{} + if s.config.Temperature != nil { + params.Temperature = s.config.Temperature + } + if s.config.MaxTokens != nil { + params.MaxTokens = s.config.MaxTokens + } + params.ToolChoice = &schemas.ToolChoice{ + ToolChoiceStr: stringPtr("auto"), + } + + // Create request + request := &schemas.BifrostRequest{ + Provider: s.config.Provider, + Model: s.config.Model, + Input: schemas.RequestInput{ + ChatCompletionInput: &s.history, + }, + Params: params, + } + + // Start loading animation + stopChan, wg := startLoader() + + // Send request + response, err := s.client.ChatCompletionRequest(context.Background(), request) + + // Stop loading animation + stopLoader(stopChan, wg) + + if err != nil { + return "", fmt.Errorf("chat completion failed: %s", err.Error.Message) + } + + if response == nil || len(response.Choices) == 0 { + return "", fmt.Errorf("no response received") + } + + // Check if response contains pending tools (requiring user approval) + if response.ExtraFields.PendingMCPTools != nil && len(*response.ExtraFields.PendingMCPTools) > 0 { + return s.handlePendingTools(response) + } + + // Get the assistant's response + choice := response.Choices[0] + assistantMessage := choice.Message + + // Add assistant message to history + s.history = append(s.history, assistantMessage) + + // Extract text content + var responseText string + if assistantMessage.Content.ContentStr != nil { + responseText = *assistantMessage.Content.ContentStr + } else if assistantMessage.Content.ContentBlocks != nil { + var textParts []string + for _, block := range *assistantMessage.Content.ContentBlocks { + if block.Text != nil { + textParts = append(textParts, *block.Text) + } + } + responseText = strings.Join(textParts, "\n") + } + + return responseText, nil +} + +// handlePendingTools handles the safe mode flow by showing pending tools and asking for approval +func (s *ChatSession) handlePendingTools(response *schemas.BifrostResponse) (string, error) { + pendingTools := *response.ExtraFields.PendingMCPTools + + // Store the assistant message with tool calls but DON'T add to history yet + // We'll add the final synthesized response instead + choice := response.Choices[0] + assistantMessage := choice.Message + + // Display pending tools to user + fmt.Println("\nπŸ”’ Safe Mode: The following tools require your approval:") + fmt.Println("=====================================================") + + for i, tool := range pendingTools { + fmt.Printf("[%d] Tool: %s\n", i+1, tool.Tool.Function.Name) + fmt.Printf(" Client: %s\n", tool.ClientName) + fmt.Printf(" Arguments: %+v\n", tool.ToolCall.Function.Arguments) + fmt.Println() + } + + fmt.Print("Do you want to approve these tools? (y/n): ") + + scanner := bufio.NewScanner(os.Stdin) + if !scanner.Scan() { + return "❌ No input received. Tool execution cancelled.", nil + } + + input := strings.ToLower(strings.TrimSpace(scanner.Text())) + if input != "y" && input != "yes" { + return "❌ Tool execution cancelled by user.", nil + } + + // First, create tool response messages for the approved tools + // This ensures proper conversation flow + toolResponseMessages := make([]schemas.BifrostMessage, 0) + + for _, pendingTool := range pendingTools { + // Create tool response message placeholder + // The actual execution will happen in the MCP plugin, but we need these in history + toolMsg := schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleTool, + Content: schemas.MessageContent{ + ContentStr: stringPtr("Tool execution approved - executing " + pendingTool.Tool.Function.Name), + }, + ToolMessage: &schemas.ToolMessage{ + ToolCallID: pendingTool.ToolCall.ID, + }, + } + toolResponseMessages = append(toolResponseMessages, toolMsg) + } + + // Note: We don't add tool response messages to persistent history here + // The final synthesized response will be added instead + + // Create approved tools list using just the tool names + // The MCP plugin will match these by name to execute the approved tools + approvedTools := make([]schemas.Tool, 0) + for _, pendingTool := range pendingTools { + approvedTools = append(approvedTools, pendingTool.Tool) + } + + // Create conversation history for approved request + // Include the assistant message with tool calls, but don't add it to our persistent history yet + conversationForApproval := append(s.history, assistantMessage) + conversationForApproval = append(conversationForApproval, toolResponseMessages...) + + // Create new request with approved tools + approvedRequest := &schemas.BifrostRequest{ + Provider: s.config.Provider, + Model: s.config.Model, + Input: schemas.RequestInput{ + ChatCompletionInput: &conversationForApproval, + }, + MCPTools: &approvedTools, + } + + fmt.Println("βœ… Tools approved") + + // Start loading animation for execution + stopChan, wg := startLoader() + + // Send approved request + approvedResponse, err := s.client.ChatCompletionRequest(context.Background(), approvedRequest) + + // Stop loading animation + stopLoader(stopChan, wg) + + if err != nil { + return "", fmt.Errorf("approved tool execution failed: %s", err.Error.Message) + } + + if approvedResponse == nil || len(approvedResponse.Choices) == 0 { + return "", fmt.Errorf("no response received from approved execution") + } + + // Get the final response + finalChoice := approvedResponse.Choices[0] + finalMessage := finalChoice.Message + + // Replace the placeholder tool messages with the actual response + // In agentic mode, this will be a synthesized response + // In non-agentic mode, this will be the tool result + if finalMessage.Role == schemas.ModelChatMessageRoleAssistant { + // This is a synthesized response from agentic mode + s.history = append(s.history, finalMessage) + } else { + // This might be a tool message in non-agentic mode, replace the last placeholder + if len(s.history) > 0 && s.history[len(s.history)-1].Role == schemas.ModelChatMessageRoleTool { + s.history[len(s.history)-1] = finalMessage + } else { + s.history = append(s.history, finalMessage) + } + } + + // Extract text content + var responseText string + if finalMessage.Content.ContentStr != nil { + responseText = *finalMessage.Content.ContentStr + } else if finalMessage.Content.ContentBlocks != nil { + var textParts []string + for _, block := range *finalMessage.Content.ContentBlocks { + if block.Text != nil { + textParts = append(textParts, *block.Text) + } + } + responseText = strings.Join(textParts, "\n") + } + + return responseText, nil +} + +// PrintHistory prints the conversation history +func (s *ChatSession) PrintHistory() { + fmt.Println("\nπŸ“œ Conversation History:") + fmt.Println("========================") + + for i, msg := range s.history { + if msg.Role == schemas.ModelChatMessageRoleSystem { + continue // Skip system messages in history display + } + + var content string + if msg.Content.ContentStr != nil { + content = *msg.Content.ContentStr + } else if msg.Content.ContentBlocks != nil { + var textParts []string + for _, block := range *msg.Content.ContentBlocks { + if block.Text != nil { + textParts = append(textParts, *block.Text) + } + } + content = strings.Join(textParts, "\n") + } + + role := strings.Title(string(msg.Role)) + timestamp := fmt.Sprintf("[%d]", i) + + fmt.Printf("%s %s: %s\n\n", timestamp, role, content) + } +} + +// Cleanup closes the chat session and cleans up resources +func (s *ChatSession) Cleanup() { + if s.client != nil { + s.client.Cleanup() + } + if s.mcpPlugin != nil { + s.mcpPlugin.Cleanup() + } +} + +// printWelcome prints the welcome message and instructions +func printWelcome(config ChatbotConfig) { + fmt.Println("πŸ€– Bifrost CLI Chatbot") + fmt.Println("======================") + fmt.Printf("πŸ”§ Provider: %s\n", config.Provider) + fmt.Printf("🧠 Model: %s\n", config.Model) + fmt.Printf("πŸ”„ Agentic Mode: %t\n", config.MCPAgenticMode) + fmt.Printf("πŸ”’ Tool Execution: Client-level policies (secure by default)\n") + if config.EnableMaximMCP { + fmt.Println("πŸ› οΈ Maxim MCP tools enabled") + } + fmt.Println() + fmt.Println("Commands:") + fmt.Println(" /help - Show this help message") + fmt.Println(" /history - Show conversation history") + fmt.Println(" /clear - Clear conversation history") + fmt.Println(" /quit - Exit the chatbot") + fmt.Println() + fmt.Println("Type your message and press Enter to chat!") + fmt.Println("==========================================") +} + +// printHelp prints help information +func printHelp() { + fmt.Println("\nπŸ“– Help") + fmt.Println("========") + fmt.Println("Available commands:") + fmt.Println(" /help - Show this help message") + fmt.Println(" /history - Show conversation history") + fmt.Println(" /clear - Clear conversation history (keeps system prompt)") + fmt.Println(" /quit - Exit the chatbot") + fmt.Println() + fmt.Println("The chatbot has access to various tools depending on configuration:") + fmt.Println("β€’ Maxim MCP tools for data operations") + fmt.Println("β€’ DuckDuckGo search for web information") + fmt.Println("β€’ In agentic mode, tool results are automatically synthesized") + fmt.Println() +} + +// stringPtr is a helper function to create string pointers +func stringPtr(s string) *string { + return &s +} + +// startLoader starts a loading spinner animation +func startLoader() (chan bool, *sync.WaitGroup) { + stopChan := make(chan bool) + var wg sync.WaitGroup + + wg.Add(1) + go func() { + defer wg.Done() + spinner := []string{"β ‹", "β ™", "β Ή", "β Έ", "β Ό", "β ΄", "β ¦", "β §", "β ‡", "⠏"} + i := 0 + + for { + select { + case <-stopChan: + // Clear the spinner + fmt.Print("\r\033[K") // Clear current line + return + default: + fmt.Printf("\rπŸ€– Assistant: %s Thinking...", spinner[i%len(spinner)]) + i++ + time.Sleep(100 * time.Millisecond) + } + } + }() + + return stopChan, &wg +} + +// stopLoader stops the loading animation +func stopLoader(stopChan chan bool, wg *sync.WaitGroup) { + close(stopChan) + wg.Wait() +} + +func main() { + // Check for required environment variables + if os.Getenv("OPENAI_API_KEY") == "" { + fmt.Println("❌ Error: OPENAI_API_KEY environment variable is required") + os.Exit(1) + } + + // Default configuration + config := ChatbotConfig{ + Provider: schemas.OpenAI, + Model: "gpt-4o-mini", + MCPAgenticMode: true, + MCPServerPort: ":8585", + EnableMaximMCP: true, + Temperature: bifrost.Ptr(0.7), + MaxTokens: bifrost.Ptr(1000), + } + + // Create chat session + fmt.Println("πŸš€ Starting Bifrost CLI Chatbot...") + session, err := NewChatSession(config) + if err != nil { + fmt.Printf("❌ Failed to create chat session: %v\n", err) + os.Exit(1) + } + + // Setup graceful shutdown + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) + go func() { + <-sigChan + fmt.Println("\n\nπŸ‘‹ Goodbye! Cleaning up...") + session.Cleanup() + os.Exit(0) + }() + + // Give MCP servers time to start + fmt.Println("⏳ Waiting for MCP servers to initialize...") + time.Sleep(3 * time.Second) + + // Print welcome message + printWelcome(config) + + // Main chat loop + scanner := bufio.NewScanner(os.Stdin) + for { + fmt.Print("\nπŸ’¬ You: ") + if !scanner.Scan() { + break + } + + input := strings.TrimSpace(scanner.Text()) + if input == "" { + continue + } + + // Handle commands + switch input { + case "/help": + printHelp() + continue + case "/history": + session.PrintHistory() + continue + case "/clear": + // Keep system prompt but clear conversation history + systemPrompt := session.history[0] // Assuming first message is system + session.history = []schemas.BifrostMessage{systemPrompt} + fmt.Println("🧹 Conversation history cleared!") + continue + case "/quit": + fmt.Println("πŸ‘‹ Goodbye!") + session.Cleanup() + return + } + + // Send message and get response + response, err := session.SendMessage(input) + if err != nil { + fmt.Printf("\rπŸ€– Assistant: ❌ Error: %v\n", err) + continue + } + + fmt.Printf("πŸ€– Assistant: %s\n", response) + } + + // Cleanup + session.Cleanup() +} diff --git a/plugins/mcp/go.mod b/plugins/mcp/go.mod new file mode 100644 index 0000000000..370ff555e8 --- /dev/null +++ b/plugins/mcp/go.mod @@ -0,0 +1,62 @@ +module github.com/maximhq/bifrost/plugins/mcp + +go 1.24.1 + +require ( + github.com/maximhq/bifrost/core v1.1.3 + github.com/metoro-io/mcp-golang v0.13.0 +) + +require ( + cloud.google.com/go/compute/metadata v0.3.0 // indirect + github.com/andybalholm/brotli v1.1.1 // indirect + github.com/aws/aws-sdk-go-v2 v1.36.3 // indirect + github.com/aws/aws-sdk-go-v2/config v1.29.14 // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.17.67 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.30 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.34 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34 // indirect + github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.25.3 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.1 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.33.19 // indirect + github.com/aws/smithy-go v1.22.3 // indirect + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/gin-contrib/sse v0.1.0 // indirect + github.com/gin-gonic/gin v1.8.1 // indirect + github.com/go-playground/locales v0.14.0 // indirect + github.com/go-playground/universal-translator v0.18.0 // indirect + github.com/go-playground/validator/v10 v10.10.0 // indirect + github.com/goccy/go-json v0.10.5 // indirect + github.com/invopop/jsonschema v0.12.0 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/compress v1.18.0 // indirect + github.com/leodido/go-urn v1.2.1 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/mattn/go-isatty v0.0.14 // indirect + github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/pelletier/go-toml/v2 v2.0.1 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/tidwall/gjson v1.18.0 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.1 // indirect + github.com/tidwall/sjson v1.2.5 // indirect + github.com/ugorji/go/codec v1.2.7 // indirect + github.com/valyala/bytebufferpool v1.0.0 // indirect + github.com/valyala/fasthttp v1.60.0 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + golang.org/x/crypto v0.37.0 // indirect + golang.org/x/net v0.39.0 // indirect + golang.org/x/oauth2 v0.30.0 // indirect + golang.org/x/sys v0.32.0 // indirect + golang.org/x/text v0.24.0 // indirect + google.golang.org/protobuf v1.28.0 // indirect + gopkg.in/yaml.v2 v2.4.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) + +replace github.com/maximhq/bifrost/core => ../../core diff --git a/plugins/mcp/go.sum b/plugins/mcp/go.sum new file mode 100644 index 0000000000..6f20369754 --- /dev/null +++ b/plugins/mcp/go.sum @@ -0,0 +1,157 @@ +cloud.google.com/go/compute/metadata v0.3.0 h1:Tz+eQXMEqDIKRsmY3cHTL6FVaynIjX2QxYC4trgAKZc= +cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k= +github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA= +github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA= +github.com/aws/aws-sdk-go-v2 v1.36.3 h1:mJoei2CxPutQVxaATCzDUjcZEjVRdpsiiXi2o38yqWM= +github.com/aws/aws-sdk-go-v2 v1.36.3/go.mod h1:LLXuLpgzEbD766Z5ECcRmi8AzSwfZItDtmABVkRLGzg= +github.com/aws/aws-sdk-go-v2/config v1.29.14 h1:f+eEi/2cKCg9pqKBoAIwRGzVb70MRKqWX4dg1BDcSJM= +github.com/aws/aws-sdk-go-v2/config v1.29.14/go.mod h1:wVPHWcIFv3WO89w0rE10gzf17ZYy+UVS1Geq8Iei34g= +github.com/aws/aws-sdk-go-v2/credentials v1.17.67 h1:9KxtdcIA/5xPNQyZRgUSpYOE6j9Bc4+D7nZua0KGYOM= +github.com/aws/aws-sdk-go-v2/credentials v1.17.67/go.mod h1:p3C44m+cfnbv763s52gCqrjaqyPikj9Sg47kUVaNZQQ= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.30 h1:x793wxmUWVDhshP8WW2mlnXuFrO4cOd3HLBroh1paFw= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.30/go.mod h1:Jpne2tDnYiFascUEs2AWHJL9Yp7A5ZVy3TNyxaAjD6M= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.34 h1:ZK5jHhnrioRkUNOc+hOgQKlUL5JeC3S6JgLxtQ+Rm0Q= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.34/go.mod h1:p4VfIceZokChbA9FzMbRGz5OV+lekcVtHlPKEO0gSZY= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34 h1:SZwFm17ZUNNg5Np0ioo/gq8Mn6u9w19Mri8DnJ15Jf0= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34/go.mod h1:dFZsC0BLo346mvKQLWmoJxT+Sjp+qcVR1tRVHQGOH9Q= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 h1:bIqFDwgGXXN1Kpp99pDOdKMTTb5d2KyU5X/BZxjOkRo= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3/go.mod h1:H5O/EsxDWyU+LP/V8i5sm8cxoZgc2fdNR9bxlOFrQTo= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3 h1:eAh2A4b5IzM/lum78bZ590jy36+d/aFLgKF/4Vd1xPE= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3/go.mod h1:0yKJC/kb8sAnmlYa6Zs3QVYqaC8ug2AbnNChv5Ox3uA= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 h1:dM9/92u2F1JbDaGooxTq18wmmFzbJRfXfVfy96/1CXM= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15/go.mod h1:SwFBy2vjtA0vZbjjaFtfN045boopadnoVPhu4Fv66vY= +github.com/aws/aws-sdk-go-v2/service/sso v1.25.3 h1:1Gw+9ajCV1jogloEv1RRnvfRFia2cL6c9cuKV2Ps+G8= +github.com/aws/aws-sdk-go-v2/service/sso v1.25.3/go.mod h1:qs4a9T5EMLl/Cajiw2TcbNt2UNo/Hqlyp+GiuG4CFDI= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.1 h1:hXmVKytPfTy5axZ+fYbR5d0cFmC3JvwLm5kM83luako= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.1/go.mod h1:MlYRNmYu/fGPoxBQVvBYr9nyr948aY/WLUvwBMBJubs= +github.com/aws/aws-sdk-go-v2/service/sts v1.33.19 h1:1XuUZ8mYJw9B6lzAkXhqHlJd/XvaX32evhproijJEZY= +github.com/aws/aws-sdk-go-v2/service/sts v1.33.19/go.mod h1:cQnB8CUnxbMU82JvlqjKR2HBOm3fe9pWorWBza6MBJ4= +github.com/aws/smithy-go v1.22.3 h1:Z//5NuZCSW6R4PhQ93hShNbyBbn8BWCmCVCt+Q8Io5k= +github.com/aws/smithy-go v1.22.3/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE= +github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI= +github.com/gin-gonic/gin v1.8.1 h1:4+fr/el88TOO3ewCmQr8cx/CtZ/umlIRIs5M4NTNjf8= +github.com/gin-gonic/gin v1.8.1/go.mod h1:ji8BvRH1azfM+SYow9zQ6SZMvR8qOMZHmsCuWR9tTTk= +github.com/go-playground/assert/v2 v2.0.1 h1:MsBgLAaY856+nPRTKrp3/OZK38U/wa0CcBYNjji3q3A= +github.com/go-playground/assert/v2 v2.0.1/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= +github.com/go-playground/locales v0.14.0 h1:u50s323jtVGugKlcYeyzC0etD1HifMjqmJqb8WugfUU= +github.com/go-playground/locales v0.14.0/go.mod h1:sawfccIbzZTqEDETgFXqTho0QybSa7l++s0DH+LDiLs= +github.com/go-playground/universal-translator v0.18.0 h1:82dyy6p4OuJq4/CByFNOn/jYrnRPArHwAcmLoJZxyho= +github.com/go-playground/universal-translator v0.18.0/go.mod h1:UvRDBj+xPUEGrFYl+lu/H90nyDXpg0fqeB/AQUGNTVA= +github.com/go-playground/validator/v10 v10.10.0 h1:I7mrTYv78z8k8VXa/qJlOlEXn/nBh+BF8dHX5nt/dr0= +github.com/go-playground/validator/v10 v10.10.0/go.mod h1:74x4gJWsvQexRdW8Pn3dXSGrTK4nAUsbPlLADvpJkos= +github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4= +github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= +github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= +github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= +github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/invopop/jsonschema v0.12.0 h1:6ovsNSuvn9wEQVOyc72aycBMVQFKz7cPdMJn10CvzRI= +github.com/invopop/jsonschema v0.12.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= +github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= +github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/leodido/go-urn v1.2.1 h1:BqpAaACuzVSgi/VLzGZIobT2z4v53pjosyNd9Yv6n/w= +github.com/leodido/go-urn v1.2.1/go.mod h1:zt4jvISO2HfUBqxjfIshjdMTYS56ZS/qv49ictyFfxY= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y= +github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= +github.com/metoro-io/mcp-golang v0.13.0 h1:54TFBJIW76VRB55CJovQQje9x4GnXg0BQQwGRtXrbCE= +github.com/metoro-io/mcp-golang v0.13.0/go.mod h1:ifLP9ZzKpN1UqFWNTpAHOqSvNkMK6b7d1FSZ5Lu0lN0= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 h1:ZqeYNhU3OHLH3mGKHDcjJRFFRrJa6eAM5H+CtDdOsPc= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/pelletier/go-toml/v2 v2.0.1 h1:8e3L2cCQzLFi2CR4g7vGFuFxX7Jl1kKX8gW+iV0GUKU= +github.com/pelletier/go-toml/v2 v2.0.1/go.mod h1:r9LEWfGN8R5k0VXJ+0BkIe7MYkRdwZOjgMj2KwnJFUo= +github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc= +github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8= +github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= +github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= +github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +github.com/ugorji/go v1.2.7/go.mod h1:nF9osbDWLy6bDVv/Rtoh6QgnvNDpmCalQV5urGCCS6M= +github.com/ugorji/go/codec v1.2.7 h1:YPXUKf7fYbp/y8xloBqZOw2qaVggbfwMlI8WM3wZUJ0= +github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY= +github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= +github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= +github.com/valyala/fasthttp v1.60.0 h1:kBRYS0lOhVJ6V+bYN8PqAHELKHtXqwq9zNMLKx1MBsw= +github.com/valyala/fasthttp v1.60.0/go.mod h1:iY4kDgV3Gc6EqhRZ8icqcmlG6bqhcDXfuHgTO4FXCvc= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= +golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= +golang.org/x/crypto v0.37.0 h1:kJNSjF/Xp7kU0iB2Z+9viTPMW4EqqsrywMXLJOOsXSE= +golang.org/x/crypto v0.37.0/go.mod h1:vg+k43peMZ0pUMhYmVAWysMK35e6ioLh3wB8ZCAfbVc= +golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= +golang.org/x/net v0.39.0 h1:ZCu7HMWDxpXpaiKdhzIfaltL9Lp31x/3fCP11bc6/fY= +golang.org/x/net v0.39.0/go.mod h1:X7NRbYVEA+ewNkCNyJ513WmMdQ3BineSwVtN2zD/d+E= +golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= +golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= +golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20210806184541-e5e7981a1069/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20= +golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0= +golang.org/x/text v0.24.0/go.mod h1:L8rBsPeo2pSS+xqN0d5u2ikmjtmoJbDBT1b7nHvFCdU= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= +google.golang.org/protobuf v1.28.0 h1:w43yiav+6bVFTBQFZX0r7ipe9JQ1QsbMgHwbBziscLw= +google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/plugins/mcp/main.go b/plugins/mcp/main.go new file mode 100644 index 0000000000..52f0b3658f --- /dev/null +++ b/plugins/mcp/main.go @@ -0,0 +1,1085 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os/exec" + "strings" + "sync" + "time" + + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + mcp_golang "github.com/metoro-io/mcp-golang" + httpTransport "github.com/metoro-io/mcp-golang/transport/http" + "github.com/metoro-io/mcp-golang/transport/stdio" +) + +// ============================================================================ +// CONSTANTS +// ============================================================================ + +const ( + // Plugin identification and defaults + PluginName = "MCPHost" // Name identifier for the MCP plugin + DefaultServerPort = ":8181" // Default port for local MCP server + BifrostVersion = "1.0.0" // Version identifier for Bifrost + BifrostClientName = "BifrostClient" // Name for internal Bifrost MCP client + BifrostClientKey = "bifrost-internal" // Key for internal Bifrost client in clientMap + LogPrefix = "[Bifrost MCP Plugin]" // Consistent logging prefix + + // Context keys for client filtering in requests + ContextKeyIncludeClients = "mcp_include_clients" // Context key for whitelist client filtering + ContextKeyExcludeClients = "mcp_exclude_clients" // Context key for blacklist client filtering +) + +// ConnectionType defines the communication protocol for MCP connections +type ConnectionType string + +const ( + ConnectionTypeHTTP ConnectionType = "http" // HTTP-based MCP connection + ConnectionTypeSTDIO ConnectionType = "stdio" // STDIO-based MCP connection +) + +// ToolExecutionPolicy defines how tools should be executed +type ToolExecutionPolicy string + +const ( + ToolExecutionPolicyRequireApproval ToolExecutionPolicy = "require_approval" // Tool requires user approval before execution + ToolExecutionPolicyAutoExecute ToolExecutionPolicy = "auto_execute" // Tool executes automatically without approval +) + +// ============================================================================ +// TYPE DEFINITIONS +// ============================================================================ + +// MCPPlugin implements schemas.Plugin for hosting and managing MCP tools. +// It provides a bridge between Bifrost and various MCP servers, supporting +// both local tool hosting and external MCP server connections. +type MCPPlugin struct { + server *mcp_golang.Server // Local MCP server instance for hosting tools + clientMap map[string]*PluginClient // Map of MCP client names to their configurations + serverPort string // Port for local MCP server + mu sync.RWMutex // Read-write mutex for thread-safe operations + agenticMode bool // Enable agentic flow (tool results sent back to LLM) + bifrostClient *bifrost.Bifrost // Bifrost client instance for agentic mode + serverRunning bool // Track whether local MCP server is running + logger schemas.Logger // Logger instance for structured logging +} + +// PluginClient represents a connected MCP client with its configuration and tools. +type PluginClient struct { + Name string // Unique name for this client + Conn *mcp_golang.Client // Active MCP client connection + ExecutionConfig ClientExecutionConfig // Tool execution policies and settings + ToolMap map[string]schemas.Tool // Available tools mapped by name + StdioCommand *exec.Cmd `json:"-"` // STDIO process command (not serialized) + ConnectionInfo ClientConnectionInfo `json:"connection_info"` // Connection metadata for management +} + +// ClientExecutionConfig defines execution policies and tool filtering for a client. +type ClientExecutionConfig struct { + Name string // Client name + DefaultPolicy ToolExecutionPolicy `json:"default_policy"` // Default execution policy for all tools + ToolPolicies map[string]ToolExecutionPolicy `json:"tool_policies,omitempty"` // Per-tool execution policies + ToolsToSkip []string `json:"tools_to_skip,omitempty"` // Tools to exclude from this client + ToolsToExecute []string // Tools to include from this client (if specified, only these are used) +} + +// ClientConnectionInfo stores metadata about how a client is connected. +type ClientConnectionInfo struct { + Type ConnectionType `json:"type"` // Connection type (HTTP or STDIO) + HTTPConnectionURL *string `json:"http_connection_url,omitempty"` // HTTP endpoint URL (for HTTP connections) + StdioCommandString *string `json:"stdio_command_string,omitempty"` // Command string for display (for STDIO connections) + ProcessID *int `json:"process_id,omitempty"` // Process ID of STDIO command +} + +// MCPPluginConfig holds configuration options for initializing the MCP plugin. +type MCPPluginConfig struct { + ServerPort string `json:"server_port,omitempty"` // Port for local MCP server (defaults to :8181) + AgenticMode bool `json:"agentic_mode,omitempty"` // Enable agentic flow for tool results + ClientConfigs []ClientExecutionConfig `json:"client_configs,omitempty"` // Per-client execution configurations +} + +// ExternalMCPConfig defines configuration for connecting to an external MCP server. +type ExternalMCPConfig struct { + Name string // Unique name for this external MCP connection + ConnectionType ConnectionType // How to connect (HTTP or STDIO) + HTTPConnectionString *string // HTTP URL (required for HTTP connections) + StdioConfig *StdioConfig // STDIO configuration (required for STDIO connections) +} + +// StdioConfig defines how to launch a STDIO-based MCP server. +type StdioConfig struct { + Command string // Executable command to run + Args []string // Command line arguments +} + +// ToolHandler is a generic function type for handling tool calls with typed arguments. +// T represents the expected argument structure for the tool. +type ToolHandler[T any] func(args T) (string, error) + +// ============================================================================ +// CONSTRUCTOR AND INITIALIZATION +// ============================================================================ + +// NewMCPPlugin creates and initializes a new MCP plugin instance. +// +// Parameters: +// - config: Plugin configuration including server port, agentic mode, and client configs +// - logger: Logger instance for structured logging (uses default if nil) +// +// Returns: +// - *MCPPlugin: Initialized plugin instance +// - error: Any initialization error +// +// The plugin will pre-create client entries for any configured clients but won't +// establish connections until ConnectToExternalMCP is called. +func NewMCPPlugin(config MCPPluginConfig, logger schemas.Logger) (*MCPPlugin, error) { + // Convert client configs to map for faster lookup during operations + clientMap := make(map[string]*PluginClient) + for _, clientConfig := range config.ClientConfigs { + clientMap[clientConfig.Name] = &PluginClient{ + Name: clientConfig.Name, + ExecutionConfig: clientConfig, + ToolMap: make(map[string]schemas.Tool), + } + } + + // Use provided logger or create default logger with info level + if logger == nil { + logger = bifrost.NewDefaultLogger(schemas.LogLevelInfo) + } + + plugin := &MCPPlugin{ + serverPort: config.ServerPort, + agenticMode: config.AgenticMode, + clientMap: clientMap, + logger: logger, + } + + plugin.logger.Info(LogPrefix + " MCP Plugin initialized") + if config.AgenticMode { + plugin.logger.Info(LogPrefix + " Agentic mode enabled") + } + + return plugin, nil +} + +// ============================================================================ +// LOCAL MCP SERVER MANAGEMENT +// ============================================================================ + +// createLocalMCPServer creates a new local MCP server instance with HTTP transport. +// This server will host tools registered via RegisterTool function. +// +// Parameters: +// - config: Plugin configuration containing server port +// +// Returns: +// - *mcp_golang.Server: Configured MCP server instance +// - error: Any creation error +func (p *MCPPlugin) createLocalMCPServer(config MCPPluginConfig) (*mcp_golang.Server, error) { + // Use configured port or default + serverPort := config.ServerPort + if serverPort == "" { + serverPort = DefaultServerPort + } + + // Create HTTP transport for the MCP server + serverTransport := httpTransport.NewHTTPTransport("/mcp") + serverTransport.WithAddr(serverPort) + server := mcp_golang.NewServer(serverTransport) + + return server, nil +} + +// createLocalMCPClient creates a client that connects to the local MCP server. +// This client is used internally by Bifrost to access locally hosted tools. +// +// Parameters: +// - config: Plugin configuration containing server port +// +// Returns: +// - *PluginClient: Configured client for local server +// - error: Any creation error +func (p *MCPPlugin) createLocalMCPClient(config MCPPluginConfig) (*PluginClient, error) { + // Use configured port or default + serverPort := config.ServerPort + if serverPort == "" { + serverPort = DefaultServerPort + } + + // Create HTTP client transport pointing to local server + clientTransport := httpTransport.NewHTTPClientTransport("/mcp") + clientTransport.WithBaseURL(fmt.Sprintf("http://localhost%s", serverPort)) + client := mcp_golang.NewClientWithInfo(clientTransport, mcp_golang.ClientInfo{ + Name: BifrostClientName, + Version: BifrostVersion, + }) + + return &PluginClient{ + Name: BifrostClientName, + Conn: client, + ExecutionConfig: ClientExecutionConfig{ + Name: BifrostClientName, + DefaultPolicy: ToolExecutionPolicyRequireApproval, + ToolPolicies: make(map[string]ToolExecutionPolicy), + }, + ToolMap: make(map[string]schemas.Tool), + }, nil +} + +// setupLocalHost initializes the local MCP server and client if not already running. +// This is called automatically when tools are registered or when the server is needed. +// +// Returns: +// - error: Any setup error +func (p *MCPPlugin) setupLocalHost() error { + // Check if server is already running + if p.server != nil && p.serverRunning { + return nil + } + + // Create and configure local MCP server + server, err := p.createLocalMCPServer(MCPPluginConfig{ServerPort: p.serverPort}) + if err != nil { + return fmt.Errorf("failed to create local MCP server: %w", err) + } + p.server = server + + // Create and configure local MCP client + client, err := p.createLocalMCPClient(MCPPluginConfig{ServerPort: p.serverPort}) + if err != nil { + return fmt.Errorf("failed to create local MCP client: %w", err) + } + p.clientMap[BifrostClientKey] = client + + // Start the server and initialize client connection + return p.startLocalMCPServer() +} + +// startLocalMCPServer starts the HTTP server and initializes the client connection. +// The server runs in a separate goroutine to avoid blocking. +// +// Returns: +// - error: Any startup error +func (p *MCPPlugin) startLocalMCPServer() error { + p.mu.Lock() + defer p.mu.Unlock() + + // Check if server is already running + if p.server != nil && p.serverRunning { + return nil + } + + if p.server == nil { + return fmt.Errorf("server not initialized") + } + + // Start the HTTP server in background goroutine + go func() { + if err := p.server.Serve(); err != nil && err != http.ErrServerClosed { + p.logger.Error(fmt.Errorf(LogPrefix+" MCP server error: %w", err)) + p.mu.Lock() + p.serverRunning = false + p.mu.Unlock() + } + }() + + // Mark server as running + p.serverRunning = true + + // Initialize the client connection to the server + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if _, ok := p.clientMap[BifrostClientKey]; !ok { + return fmt.Errorf("bifrost client not found") + } + + _, err := p.clientMap[BifrostClientKey].Conn.Initialize(ctx) + if err != nil { + p.serverRunning = false + return fmt.Errorf("failed to initialize MCP client: %v", err) + } + + return nil +} + +// ============================================================================ +// TOOL REGISTRATION +// ============================================================================ + +// RegisterTool registers a typed tool handler with the local MCP server. +// This is a convenience function that handles the conversion between typed Go +// handlers and the MCP protocol. +// +// Type Parameters: +// - T: The expected argument type for the tool (must be JSON-deserializable) +// +// Parameters: +// - plugin: The MCP plugin instance +// - name: Unique tool name +// - description: Human-readable tool description +// - handler: Typed function that handles tool execution +// - toolSchema: Bifrost tool schema for function calling +// - policy: Execution policy for this tool +// +// Returns: +// - error: Any registration error +// +// Example: +// +// type EchoArgs struct { +// Message string `json:"message"` +// } +// +// RegisterTool(plugin, "echo", "Echo a message", +// func(args EchoArgs) (string, error) { +// return args.Message, nil +// }, toolSchema, ToolExecutionPolicyAutoExecute) +func RegisterTool[T any](plugin *MCPPlugin, name, description string, handler ToolHandler[T], toolSchema schemas.Tool, policy ToolExecutionPolicy) error { + // Ensure local server is set up + if err := plugin.setupLocalHost(); err != nil { + return fmt.Errorf("failed to setup local host: %w", err) + } + + // Verify internal client exists + if _, ok := plugin.clientMap[BifrostClientKey]; !ok { + return fmt.Errorf("bifrost client not found") + } + + plugin.mu.Lock() + defer plugin.mu.Unlock() + + plugin.logger.Info(fmt.Sprintf(LogPrefix+" Registering typed tool: %s", name)) + + // Create MCP handler wrapper that converts between typed and MCP interfaces + mcpHandler := func(args T) (*mcp_golang.ToolResponse, error) { + result, err := handler(args) + if err != nil { + return mcp_golang.NewToolResponse(mcp_golang.NewTextContent(fmt.Sprintf("Error: %s", err.Error()))), nil + } + return mcp_golang.NewToolResponse(mcp_golang.NewTextContent(result)), nil + } + + // Register with the underlying mcp-golang server + err := plugin.server.RegisterTool(name, description, mcpHandler) + if err != nil { + return fmt.Errorf("failed to register tool with MCP server: %w", err) + } + + // Store tool definition and policy for Bifrost integration + plugin.clientMap[BifrostClientKey].ToolMap[name] = toolSchema + plugin.clientMap[BifrostClientKey].ExecutionConfig.ToolPolicies[name] = policy + + return nil +} + +// ============================================================================ +// EXTERNAL MCP CONNECTION MANAGEMENT +// ============================================================================ + +// ConnectToExternalMCP establishes a connection to an external MCP server and +// registers its available tools with the plugin. +// +// Supported connection types: +// - HTTP: Connects to an HTTP-based MCP server +// - STDIO: Launches and connects to a command-line MCP server +// +// Parameters: +// - config: External MCP connection configuration +// +// Returns: +// - error: Any connection or registration error +// +// The function will: +// 1. Create or update the client entry in clientMap +// 2. Establish the connection based on the specified type +// 3. Initialize the connection and retrieve available tools +// 4. Register tools with the plugin (subject to filtering rules) +func (p *MCPPlugin) ConnectToExternalMCP(config ExternalMCPConfig) error { + p.mu.Lock() + defer p.mu.Unlock() + + // Initialize or validate client entry + if existingClient, exists := p.clientMap[config.Name]; exists { + // Client entry exists from config, check for existing connection + if existingClient.Conn != nil { + return fmt.Errorf("client %s already has an active connection", config.Name) + } + // Update connection type for this connection attempt + existingClient.ConnectionInfo.Type = config.ConnectionType + } else { + // Create new client entry with default configuration + p.clientMap[config.Name] = &PluginClient{ + Name: config.Name, + ExecutionConfig: ClientExecutionConfig{ + Name: config.Name, + DefaultPolicy: ToolExecutionPolicyRequireApproval, + ToolPolicies: make(map[string]ToolExecutionPolicy), + ToolsToSkip: make([]string, 0), + }, + ToolMap: make(map[string]schemas.Tool), + ConnectionInfo: ClientConnectionInfo{ + Type: config.ConnectionType, + }, + } + } + + var externalClient *mcp_golang.Client + var err error + + // Create appropriate transport based on connection type + switch config.ConnectionType { + case ConnectionTypeHTTP: + externalClient, err = p.createHTTPConnection(config) + case ConnectionTypeSTDIO: + externalClient, err = p.createSTDIOConnection(config) + default: + return fmt.Errorf("unknown connection type: %s", config.ConnectionType) + } + + if err != nil { + return fmt.Errorf("failed to create connection: %w", err) + } + + // Initialize the external client with timeout + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + _, err = externalClient.Initialize(ctx) + if err != nil { + return fmt.Errorf("failed to initialize external MCP client %s: %v", config.Name, err) + } + + // Retrieve and register available tools from the external server + err = p.registerExternalTools(ctx, externalClient, config.Name) + if err != nil { + p.logger.Warn(fmt.Sprintf(LogPrefix+" Failed to register tools from %s: %v", config.Name, err)) + // Continue with connection even if tool registration fails + } + + // Store the external client connection + p.clientMap[config.Name].Conn = externalClient + + return nil +} + +// createHTTPConnection creates an HTTP-based MCP client connection. +func (p *MCPPlugin) createHTTPConnection(config ExternalMCPConfig) (*mcp_golang.Client, error) { + if config.HTTPConnectionString == nil { + return nil, fmt.Errorf("HTTP connection string is required") + } + + // Store HTTP connection info + p.clientMap[config.Name].ConnectionInfo.HTTPConnectionURL = config.HTTPConnectionString + + // Create HTTP transport + clientTransport := httpTransport.NewHTTPClientTransport("/mcp") + clientTransport.WithBaseURL(*config.HTTPConnectionString) + + return mcp_golang.NewClientWithInfo(clientTransport, mcp_golang.ClientInfo{ + Name: fmt.Sprintf("Bifrost-%s", config.Name), + Version: "1.0.0", + }), nil +} + +// createSTDIOConnection creates a STDIO-based MCP client connection. +func (p *MCPPlugin) createSTDIOConnection(config ExternalMCPConfig) (*mcp_golang.Client, error) { + if config.StdioConfig == nil { + return nil, fmt.Errorf("stdio config is required") + } + + // Store STDIO command info for display + cmdString := fmt.Sprintf("%s %s", config.StdioConfig.Command, strings.Join(config.StdioConfig.Args, " ")) + p.clientMap[config.Name].ConnectionInfo.StdioCommandString = &cmdString + + // Create and start the STDIO command + cmd := exec.Command(config.StdioConfig.Command, config.StdioConfig.Args...) + + // Get stdin/stdout pipes before starting + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, fmt.Errorf("failed to get stdin pipe: %v", err) + } + stdout, err := cmd.StdoutPipe() + if err != nil { + stdin.Close() // Clean up stdin if stdout fails + return nil, fmt.Errorf("failed to get stdout pipe: %v", err) + } + + // Start the command + if err := cmd.Start(); err != nil { + stdin.Close() + stdout.Close() + return nil, fmt.Errorf("failed to start command '%s %v': %v", config.StdioConfig.Command, config.StdioConfig.Args, err) + } + + // Track the command and process ID for cleanup + p.clientMap[config.Name].StdioCommand = cmd + if cmd.Process != nil { + pid := cmd.Process.Pid + p.clientMap[config.Name].ConnectionInfo.ProcessID = &pid + } + + // Create stdio transport with the command's stdout as our stdin, and stdin as our stdout + stdioTransport := stdio.NewStdioServerTransportWithIO(stdout, stdin) + + return mcp_golang.NewClientWithInfo(stdioTransport, mcp_golang.ClientInfo{ + Name: fmt.Sprintf("Bifrost-%s", config.Name), + Version: "1.0.0", + }), nil +} + +// registerExternalTools retrieves and registers tools from an external MCP server. +func (p *MCPPlugin) registerExternalTools(ctx context.Context, client *mcp_golang.Client, clientName string) error { + // Get available tools from external server + // Pass empty string instead of nil to avoid "Expected string, received null" error + toolsResponse, err := client.ListTools(ctx, bifrost.Ptr("")) + if err != nil { + return fmt.Errorf("failed to list tools: %v", err) + } + + if toolsResponse == nil { + return nil // No tools available + } + + // Convert and register each tool + for _, mcpTool := range toolsResponse.Tools { + // Skip tools that are configured to be skipped + if p.shouldSkipTool(mcpTool.Name, clientName) { + continue + } + + // Convert MCP tool schema to Bifrost format + bifrostTool := convertMCPToolToBifrostSchema(&mcpTool) + p.clientMap[clientName].ToolMap[mcpTool.Name] = bifrostTool + } + + return nil +} + +// ============================================================================ +// PLUGIN INTERFACE IMPLEMENTATION +// ============================================================================ + +// GetName returns the plugin's name identifier. +// This implements the schemas.Plugin interface. +func (p *MCPPlugin) GetName() string { + return PluginName +} + +// SetBifrostClient sets the Bifrost client instance for agentic mode. +// This client is used to make follow-up requests when agentic mode is enabled. +// +// Parameters: +// - client: Bifrost client instance +func (p *MCPPlugin) SetBifrostClient(client *bifrost.Bifrost) { + p.mu.Lock() + defer p.mu.Unlock() + p.bifrostClient = client +} + +// PreHook is called before request processing to add available MCP tools. +// This implements the schemas.Plugin interface. +// +// The function: +// 1. Handles approved tool execution (from user approval flow) +// 2. Adds available MCP tools to the request for normal flow +// 3. Applies client filtering based on request context +// +// Parameters: +// - ctx: Request context (may contain client filtering preferences) +// - req: Incoming Bifrost request +// +// Returns: +// - *schemas.BifrostRequest: Modified request with MCP tools +// - *schemas.BifrostResponse: Response if tools were executed (short-circuit) +// - error: Any processing error +func (p *MCPPlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.BifrostResponse, error) { + // Check if this is an approved tool execution request + if req.MCPTools != nil && len(*req.MCPTools) > 0 && req.Input.ChatCompletionInput != nil { + return p.handleApprovedTools(ctx, req) + } + + // Normal flow: Add available tools to request + availableTools := p.getFilteredAvailableTools(ctx) + + // Initialize tools array if needed + if req.Params == nil { + req.Params = &schemas.ModelParameters{} + } + if req.Params.Tools == nil { + req.Params.Tools = &[]schemas.Tool{} + } + tools := *req.Params.Tools + + // Add MCP tools, avoiding duplicates + for _, mcpTool := range availableTools { + isDuplicate := false + for _, tool := range tools { + if tool.Function.Name == mcpTool.Function.Name { + isDuplicate = true + break + } + } + if !isDuplicate { + tools = append(tools, mcpTool) + } + } + + req.Params.Tools = &tools + return req, nil, nil +} + +// PostHook is called after response generation to handle tool calls. +// This implements the schemas.Plugin interface. +// +// The function: +// 1. Detects tool calls in the response +// 2. Applies execution policies (auto-execute vs require approval) +// 3. Executes approved tools or returns pending tools for user approval +// 4. Handles agentic flow if enabled +// +// Parameters: +// - ctx: Request context +// - res: Generated response from LLM +// - err: Any error from response generation +// +// Returns: +// - *schemas.BifrostResponse: Modified response with tool results or pending tools +// - *schemas.BifrostError: Any processing error +// - error: Any fatal error +func (p *MCPPlugin) PostHook(ctx *context.Context, res *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) { + if res == nil || res.Choices == nil { + return res, err, nil + } + + // Check each choice for tool calls + for i, choice := range res.Choices { + if choice.Message.ToolCalls != nil && len(*choice.Message.ToolCalls) > 0 { + return p.handleToolCallsWithPolicy(ctx, res, i, choice) + } + } + + return res, err, nil +} + +// Cleanup performs cleanup of all resources when the plugin is being destroyed. +// This implements the schemas.Plugin interface. +// +// The function: +// 1. Terminates all STDIO processes +// 2. Disconnects all MCP clients +// 3. Clears server references +// +// Returns: +// - error: Any cleanup error +func (p *MCPPlugin) Cleanup() error { + p.mu.Lock() + defer p.mu.Unlock() + + // Clean up STDIO processes + for _, client := range p.clientMap { + if client.StdioCommand != nil && client.StdioCommand.Process != nil { + p.logger.Info(fmt.Sprintf(LogPrefix+" Terminating STDIO process: %d", client.StdioCommand.Process.Pid)) + client.StdioCommand.Process.Kill() + client.StdioCommand.Wait() // Wait for cleanup + } + } + + // Disconnect all clients + for name := range p.clientMap { + p.logger.Info(fmt.Sprintf(LogPrefix+" Disconnecting MCP client: %s", name)) + } + p.clientMap = make(map[string]*PluginClient) + + // Clear server reference + if p.server != nil { + p.logger.Info(LogPrefix + " Clearing MCP server reference") + p.server = nil + p.serverRunning = false + } + + return nil +} + +// ============================================================================ +// TOOL EXECUTION AND FLOW MANAGEMENT +// ============================================================================ + +// getFilteredAvailableTools returns tools filtered by request-level client inclusion/exclusion. +// Client filtering allows requests to specify which MCP clients' tools should be included. +// +// Parameters: +// - ctx: Request context containing potential filtering directives +// +// Returns: +// - []schemas.Tool: List of available tools after applying filters +func (p *MCPPlugin) getFilteredAvailableTools(ctx *context.Context) []schemas.Tool { + p.mu.RLock() + defer p.mu.RUnlock() + + var includeClients []string + var excludeClients []string + + // Extract client filtering from request context + if ctx != nil { + if existingIncludeClients, ok := (*ctx).Value(ContextKeyIncludeClients).([]string); ok && existingIncludeClients != nil { + includeClients = existingIncludeClients + } + if existingExcludeClients, ok := (*ctx).Value(ContextKeyExcludeClients).([]string); ok && existingExcludeClients != nil { + excludeClients = existingExcludeClients + } + } + + tools := make([]schemas.Tool, 0) + for clientName, client := range p.clientMap { + // Apply client filtering logic + if !p.shouldIncludeClient(clientName, includeClients, excludeClients) { + continue + } + + // Add all tools from this client + for _, tool := range client.ToolMap { + tools = append(tools, tool) + } + } + return tools +} + +// callTool executes a tool by finding the appropriate MCP client and invoking the tool. +// +// Parameters: +// - ctx: Execution context +// - toolName: Name of the tool to execute +// - arguments: Tool arguments as key-value pairs +// +// Returns: +// - *mcp_golang.ToolResponse: Tool execution result +// - error: Any execution error +func (p *MCPPlugin) callTool(ctx context.Context, toolName string, arguments map[string]interface{}) (*mcp_golang.ToolResponse, error) { + // Find which client has this tool + client := p.findMCPClientForTool(toolName) + if client == nil { + return nil, fmt.Errorf("tool '%s' not found in any connected MCP client", toolName) + } + + if client.Conn == nil { + return nil, fmt.Errorf("client '%s' has no active connection", client.Name) + } + + return client.Conn.CallTool(ctx, toolName, arguments) +} + +// handleApprovedTools executes MCP tools that have been approved by the user. +// This is called when a request contains pre-approved tools. +// +// Parameters: +// - ctx: Request context +// - req: Request containing approved tools +// +// Returns: +// - *schemas.BifrostRequest: Modified request for agentic flow +// - *schemas.BifrostResponse: Response with tool results (non-agentic) +// - error: Any execution error +func (p *MCPPlugin) handleApprovedTools(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.BifrostResponse, error) { + // Validate request has conversation history + if req.Input.ChatCompletionInput == nil || len(*req.Input.ChatCompletionInput) == 0 { + return req, nil, nil + } + + messages := *req.Input.ChatCompletionInput + + // Find the assistant message with tool calls + var assistantMessageWithToolCalls *schemas.BifrostMessage + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role == schemas.ModelChatMessageRoleAssistant && messages[i].ToolCalls != nil && len(*messages[i].ToolCalls) > 0 { + assistantMessageWithToolCalls = &messages[i] + break + } + } + + if assistantMessageWithToolCalls == nil { + return req, nil, nil + } + + // Create approved tools lookup map + approvedTools := make(map[string]schemas.Tool) + for _, tool := range *req.MCPTools { + approvedTools[tool.Function.Name] = tool + } + + toolCallResults := make([]schemas.BifrostMessage, 0) + + // Execute each approved tool + for _, toolCall := range *assistantMessageWithToolCalls.ToolCalls { + if toolCall.Function.Name == nil { + continue + } + toolName := *toolCall.Function.Name + + // Verify tool is approved and is an MCP tool + if _, isApproved := approvedTools[toolName]; !isApproved { + continue + } + + client := p.findMCPClientForTool(toolName) + if client == nil { + continue + } + + // Execute the tool + toolMsg, err := p.executeSingleTool(context.Background(), toolCall) + if err != nil { + return nil, &schemas.BifrostResponse{}, err + } + + toolCallResults = append(toolCallResults, toolMsg) + } + + // Handle results based on agentic mode + if len(toolCallResults) == 0 { + return req, nil, nil + } + + if p.checkAgenticModeAvailable() { + // Agentic mode: Add tool results to conversation and continue to LLM + return p.prepareAgenticRequest(req, toolCallResults), nil, nil + } + + // Non-agentic mode: Return tool results directly + response := &schemas.BifrostResponse{ + Choices: []schemas.BifrostResponseChoice{ + { + Index: 0, + Message: toolCallResults[0], // Return first tool result for backwards compatibility + }, + }, + } + return nil, response, nil +} + +// prepareAgenticRequest prepares a request for agentic flow by adding tool results to conversation. +func (p *MCPPlugin) prepareAgenticRequest(req *schemas.BifrostRequest, toolCallResults []schemas.BifrostMessage) *schemas.BifrostRequest { + conversationHistory := *req.Input.ChatCompletionInput + + // Remove placeholder tool messages to avoid duplicates + cleanedHistory := make([]schemas.BifrostMessage, 0) + for _, msg := range conversationHistory { + // Skip placeholder tool messages + if msg.Role == schemas.ModelChatMessageRoleTool && + msg.Content.ContentStr != nil && + strings.Contains(*msg.Content.ContentStr, "Tool execution approved") { + continue + } + cleanedHistory = append(cleanedHistory, msg) + } + + // Add actual tool results + cleanedHistory = append(cleanedHistory, toolCallResults...) + + // Add synthesis prompt + synthesisPrompt := schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ + ContentStr: bifrost.Ptr("Please provide a comprehensive response based on the tool results above."), + }, + } + cleanedHistory = append(cleanedHistory, synthesisPrompt) + + // Update request + req.Input.ChatCompletionInput = &cleanedHistory + req.MCPTools = nil // Clear MCP tools since we're not using them in the next turn + + return req +} + +// handleToolCallsWithPolicy processes tool calls based on their execution policies. +// Tools are categorized as either requiring approval or auto-executing. +// +// Parameters: +// - ctx: Request context +// - res: Response containing tool calls +// - choiceIndex: Index of the choice being processed +// - choice: The specific choice containing tool calls +// +// Returns: +// - *schemas.BifrostResponse: Modified response +// - *schemas.BifrostError: Any processing error +// - error: Any fatal error +func (p *MCPPlugin) handleToolCallsWithPolicy(ctx *context.Context, res *schemas.BifrostResponse, choiceIndex int, choice schemas.BifrostResponseChoice) (*schemas.BifrostResponse, *schemas.BifrostError, error) { + pendingTools := make([]schemas.PendingMCPTool, 0) + autoExecuteTools := make([]schemas.ToolCall, 0) + + // Categorize tools based on execution policies + for _, toolCall := range *choice.Message.ToolCalls { + if toolCall.Function.Name == nil { + continue + } + toolName := *toolCall.Function.Name + + client := p.findMCPClientForTool(toolName) + if client == nil { + continue // Skip tools not found in any MCP client + } + + if p.shouldRequireApproval(toolName, client.Name) { + // Tool requires user approval + var arguments map[string]interface{} + if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &arguments); err != nil { + return nil, &schemas.BifrostError{ + Error: schemas.ErrorField{ + Message: fmt.Sprintf("Failed to parse tool arguments: %v", err), + }, + }, nil + } + + pendingTool := schemas.PendingMCPTool{ + ClientName: client.Name, + Tool: client.ToolMap[toolName], + ToolCall: toolCall, + } + pendingTools = append(pendingTools, pendingTool) + } else { + // Tool can be auto-executed + autoExecuteTools = append(autoExecuteTools, toolCall) + } + } + + // Handle pending tools (require user approval) + if len(pendingTools) > 0 { + res.ExtraFields.PendingMCPTools = &pendingTools + return res, nil, nil + } + + // Handle auto-execute tools + if len(autoExecuteTools) > 0 { + return p.executeToolsImmediately(ctx, res, choiceIndex, choice, autoExecuteTools) + } + + return res, nil, nil +} + +// executeToolsImmediately executes tools that are configured for automatic execution. +// +// Parameters: +// - ctx: Request context +// - res: Response to modify +// - choiceIndex: Index of choice being processed +// - choice: Choice containing tool calls +// - toolsToExecute: List of tools to execute +// +// Returns: +// - *schemas.BifrostResponse: Modified response +// - *schemas.BifrostError: Any processing error +// - error: Any fatal error +func (p *MCPPlugin) executeToolsImmediately(ctx *context.Context, res *schemas.BifrostResponse, choiceIndex int, choice schemas.BifrostResponseChoice, toolsToExecute []schemas.ToolCall) (*schemas.BifrostResponse, *schemas.BifrostError, error) { + toolCallResults := make([]schemas.BifrostMessage, 0) + assistantMessage := choice.Message // Preserve original assistant message + + // Execute each tool + for _, toolCall := range toolsToExecute { + toolMsg, err := p.executeSingleTool(*ctx, toolCall) + if err != nil { + return nil, &schemas.BifrostError{ + Error: schemas.ErrorField{ + Message: err.Error(), + }, + }, nil + } + toolCallResults = append(toolCallResults, toolMsg) + } + + // Handle results based on agentic mode + if len(toolCallResults) > 0 { + if p.checkAgenticModeAvailable() { + // Agentic mode: Send conversation back to LLM for synthesis + return p.handleAgenticFlow(ctx, res, choiceIndex, assistantMessage, toolCallResults) + } else { + // Non-agentic mode: Replace with tool result + res.Choices[choiceIndex].Message = toolCallResults[0] + } + } + + return res, nil, nil +} + +// handleAgenticFlow processes tool results in agentic mode by sending the conversation +// back to the LLM for synthesis and natural language response generation. +// +// Parameters: +// - ctx: Request context +// - res: Original response +// - choiceIndex: Index of choice being processed +// - assistantMessage: Original assistant message with tool calls +// - toolCallResults: Results from tool execution +// +// Returns: +// - *schemas.BifrostResponse: Response with synthesized content +// - *schemas.BifrostError: Any processing error +// - error: Any fatal error +func (p *MCPPlugin) handleAgenticFlow(ctx *context.Context, res *schemas.BifrostResponse, choiceIndex int, assistantMessage schemas.BifrostMessage, toolCallResults []schemas.BifrostMessage) (*schemas.BifrostResponse, *schemas.BifrostError, error) { + // Verify agentic mode is properly configured + if !p.checkAgenticModeAvailable() { + // Fallback to non-agentic mode + if len(toolCallResults) > 0 { + res.Choices[choiceIndex].Message = toolCallResults[0] + } + return res, nil, nil + } + + // Reconstruct conversation history + var conversationHistory []schemas.BifrostMessage + if res.ExtraFields.ChatHistory != nil { + conversationHistory = *res.ExtraFields.ChatHistory + } + + // Add assistant message with tool calls + conversationHistory = append(conversationHistory, assistantMessage) + + // Add all tool results + conversationHistory = append(conversationHistory, toolCallResults...) + + // Add synthesis prompt + synthesisPrompt := schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ + ContentStr: bifrost.Ptr("Please provide a comprehensive response based on the tool results above."), + }, + } + conversationHistory = append(conversationHistory, synthesisPrompt) + + // Create agentic request + agenticRequest := &schemas.BifrostRequest{ + Provider: res.ExtraFields.Provider, + Model: "gpt-4o-mini", + Input: schemas.RequestInput{ + ChatCompletionInput: &conversationHistory, + }, + Params: &res.ExtraFields.Params, + } + + // Make agentic call + agenticResponse, bifrostErr := p.bifrostClient.ChatCompletionRequest(context.Background(), agenticRequest) + if bifrostErr != nil { + p.logger.Warn(fmt.Sprintf(LogPrefix+" Agentic call failed: %v. Falling back to normal execution.", bifrostErr.Error.Message)) + // Fallback to non-agentic mode + if len(toolCallResults) > 0 { + res.Choices[choiceIndex].Message = toolCallResults[0] + } + return res, nil, nil + } + + // Replace original choice with synthesized response + if agenticResponse != nil && len(agenticResponse.Choices) > 0 { + res.Choices[choiceIndex] = agenticResponse.Choices[0] + res.ExtraFields.ChatHistory = &conversationHistory + } + + return res, nil, nil +} diff --git a/plugins/mcp/plugin_test.go b/plugins/mcp/plugin_test.go new file mode 100644 index 0000000000..8d2560f325 --- /dev/null +++ b/plugins/mcp/plugin_test.go @@ -0,0 +1,684 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + "testing" + "time" + + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" +) + +// WeatherArgs defines the arguments for the weather tool +type WeatherArgs struct { + City string `json:"city" jsonschema:"required,description=The city name to get weather for"` +} + +// Mock weather data for testing +var mockWeatherData = map[string]string{ + "new york": "Sunny, 22Β°C", + "london": "Cloudy, 15Β°C", + "tokyo": "Rainy, 18Β°C", + "paris": "Partly cloudy, 19Β°C", + "san francisco": "Foggy, 16Β°C", +} + +// BaseAccount implements the schemas.Account interface for testing purposes. +type BaseAccount struct{} + +func (baseAccount *BaseAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) { + return []schemas.ModelProvider{schemas.OpenAI}, nil +} + +func (baseAccount *BaseAccount) GetKeysForProvider(providerKey schemas.ModelProvider) ([]schemas.Key, error) { + return []schemas.Key{ + { + Value: os.Getenv("OPENAI_API_KEY"), + Models: []string{"gpt-4o-mini", "gpt-4-turbo"}, + Weight: 1.0, + }, + }, nil +} + +func (baseAccount *BaseAccount) GetConfigForProvider(providerKey schemas.ModelProvider) (*schemas.ProviderConfig, error) { + return &schemas.ProviderConfig{ + NetworkConfig: schemas.DefaultNetworkConfig, + ConcurrencyAndBufferSize: schemas.DefaultConcurrencyAndBufferSize, + }, nil +} + +func TestMCPPlugin_WeatherTool(t *testing.T) { + // Check if OpenAI API key is available + if os.Getenv("OPENAI_API_KEY") == "" { + t.Skip("OPENAI_API_KEY not set, skipping integration test") + } + + // Create the MCP host plugin + plugin, err := NewMCPPlugin(MCPPluginConfig{ServerPort: ":8282"}, bifrost.NewDefaultLogger(schemas.LogLevelDebug)) // Use a different port + if err != nil { + t.Fatalf("Failed to create MCP plugin: %v", err) + } + + // Define the weather tool schema for Bifrost + weatherToolSchema := schemas.Tool{ + Type: "function", + Function: schemas.Function{ + Name: "get_weather", + Description: "Get current weather information for a specified city", + Parameters: schemas.FunctionParameters{ + Type: "object", + Properties: map[string]interface{}{ + "city": map[string]interface{}{ + "type": "string", + "description": "The city name to get weather for", + }, + }, + Required: []string{"city"}, + }, + }, + } + + // Register weather tool using plugin's API + weatherHandler := func(args WeatherArgs) (string, error) { + // Case-insensitive lookup + cityLower := strings.ToLower(args.City) + if weather, exists := mockWeatherData[cityLower]; exists { + return fmt.Sprintf("Weather in %s: %s", args.City, weather), nil + } + return fmt.Sprintf("Weather data not available for %s", args.City), nil + } + + err = RegisterTool(plugin, "get_weather", "Get current weather information for a specified city", + weatherHandler, weatherToolSchema, ToolExecutionPolicyAutoExecute) + if err != nil { + t.Fatalf("Failed to register weather tool: %v", err) + } + + // Initialize Bifrost with the MCP plugin + account := BaseAccount{} + client, err := bifrost.Init(schemas.BifrostConfig{ + Account: &account, + Plugins: []schemas.Plugin{plugin}, + Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug), + }) + if err != nil { + t.Fatalf("Failed to initialize Bifrost: %v", err) + } + + t.Run("WeatherQuery_ThroughBifrost", func(t *testing.T) { + fmt.Println("=== BIFROST WEATHER QUERY TEST ===") + + // Make a chat completion request that should trigger the weather tool + response, bifrostErr := client.ChatCompletionRequest(context.Background(), &schemas.BifrostRequest{ + Provider: schemas.OpenAI, + Model: "gpt-4o-mini", + Input: schemas.RequestInput{ + ChatCompletionInput: &[]schemas.BifrostMessage{ + { + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ContentStr: stringPtr("What's the weather like in New York? Use the get_weather tool.")}, + }, + }, + }, + }) + + if bifrostErr != nil { + t.Fatalf("Chat completion request failed: %v", bifrostErr) + } + + if response == nil { + t.Fatalf("Expected response, got nil") + } + + fmt.Printf("Response received with %d choices\n", len(response.Choices)) + + // Check if we got a response + if len(response.Choices) == 0 { + t.Fatalf("Expected at least one choice in response") + } + + // Log the response for debugging + choice := response.Choices[0] + fmt.Printf("Response role: %s\n", choice.Message.Role) + + if choice.Message.Content.ContentStr != nil { + fmt.Printf("Response content: %s\n", *choice.Message.Content.ContentStr) + + // Check if the response contains weather information + responseText := strings.ToLower(*choice.Message.Content.ContentStr) + if !strings.Contains(responseText, "weather") && !strings.Contains(responseText, "sunny") && + !strings.Contains(responseText, "cloudy") && !strings.Contains(responseText, "Β°c") { + t.Logf("Warning: Response may not contain expected weather information") + } + } + + // Check if tool calls were made (could be in the response) + if choice.Message.AssistantMessage != nil && choice.Message.AssistantMessage.ToolCalls != nil { + fmt.Printf("Response contains %d tool calls\n", len(*choice.Message.AssistantMessage.ToolCalls)) + for i, toolCall := range *choice.Message.AssistantMessage.ToolCalls { + fmt.Printf(" Tool call %d: %s\n", i+1, *toolCall.Function.Name) + } + } + }) + + // Cleanup + client.Cleanup() + fmt.Println("\n=== TEST COMPLETED SUCCESSFULLY ===") +} + +// Helper function to print choices as JSON +func printChoicesAsJSON(t *testing.T, choices []schemas.BifrostResponseChoice) { + t.Helper() + jsonData, err := json.MarshalIndent(choices, "", " ") + if err != nil { + t.Errorf("Failed to marshal choices to JSON: %v", err) + return + } + fmt.Printf("--- Response Choices (JSON) ---\n%s\n-----------------------------\n", string(jsonData)) +} + +// Helper function to create string pointer +func stringPtr(s string) *string { + return &s +} + +func TestMCPPlugin_Context7Integration(t *testing.T) { + // Skip this test if no OPENAI_API_KEY is set + if os.Getenv("OPENAI_API_KEY") == "" { + t.Skip("Skipping Context7 integration test: OPENAI_API_KEY not set") + } + + // Create MCP plugin + plugin, err := NewMCPPlugin(MCPPluginConfig{}, bifrost.NewDefaultLogger(schemas.LogLevelDebug)) + if err != nil { + t.Fatalf("Failed to create plugin: %v", err) + } + + // Connect to Context7 MCP server using npx + err = plugin.ConnectToExternalMCP(ExternalMCPConfig{ + Name: "context7", + ConnectionType: ConnectionTypeSTDIO, + StdioConfig: &StdioConfig{ + Command: "npx", + Args: []string{"-y", "@upstash/context7-mcp"}, + }, + }) + if err != nil { + t.Fatalf("Failed to connect to Context7 MCP: %v", err) + } + + // Give some time for the external server to start and register tools + time.Sleep(2 * time.Second) + + // Initialize Bifrost with the MCP plugin + account := BaseAccount{} + client, err := bifrost.Init(schemas.BifrostConfig{ + Account: &account, + Plugins: []schemas.Plugin{plugin}, + Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug), + }) + if err != nil { + t.Fatalf("Failed to initialize Bifrost: %v", err) + } + + // Test 1: Check if Context7 tools are available through the plugin + availableTools := plugin.getFilteredAvailableTools(nil) + + expectedTools := []string{"resolve-library-id", "get-library-docs"} + for _, expectedTool := range expectedTools { + found := false + for _, tool := range availableTools { + if tool.Function.Name == expectedTool { + found = true + break + } + } + if !found { + t.Errorf("Expected Context7 tool '%s' not found in available tools", expectedTool) + } + } + + // Test 2: Use Context7 to resolve a library ID through Bifrost + t.Run("ResolveLibraryID_ThroughBifrost", func(t *testing.T) { + fmt.Println("=== CONTEXT7 RESOLVE LIBRARY ID TEST ===") + + response, bifrostErr := client.ChatCompletionRequest(context.Background(), &schemas.BifrostRequest{ + Provider: schemas.OpenAI, + Model: "gpt-4o-mini", + Input: schemas.RequestInput{ + ChatCompletionInput: &[]schemas.BifrostMessage{ + { + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ContentStr: stringPtr("I need to resolve the library ID for 'react' using the resolve-library-id tool.")}, + }, + }, + }, + Params: &schemas.ModelParameters{ + ToolChoice: &schemas.ToolChoice{ + ToolChoiceStr: stringPtr("auto"), + }, + }, + }) + + if bifrostErr != nil { + t.Fatalf("Chat completion request failed: %v", bifrostErr) + } + + if response == nil || len(response.Choices) == 0 { + t.Fatal("Expected response with choices") + } + + choice := response.Choices[0] + fmt.Printf("Response role: %s\n", choice.Message.Role) + + if choice.Message.AssistantMessage != nil && choice.Message.AssistantMessage.ToolCalls != nil { + fmt.Printf("Response contains %d tool calls\n", len(*choice.Message.AssistantMessage.ToolCalls)) + for i, toolCall := range *choice.Message.AssistantMessage.ToolCalls { + fmt.Printf(" Tool call %d: %s\n", i+1, *toolCall.Function.Name) + if *toolCall.Function.Name == "resolve-library-id" { + t.Logf("Successfully called resolve-library-id with arguments: %s", toolCall.Function.Arguments) + } + } + } else { + t.Log("No tool calls made - this might be expected if the model chooses not to call tools") + } + }) + + // Test 3: Use Context7 to get documentation + t.Run("GetLibraryDocs_ThroughBifrost", func(t *testing.T) { + fmt.Println("=== CONTEXT7 GET LIBRARY DOCS TEST ===") + + response, bifrostErr := client.ChatCompletionRequest(context.Background(), &schemas.BifrostRequest{ + Provider: schemas.OpenAI, + Model: "gpt-4o-mini", + Input: schemas.RequestInput{ + ChatCompletionInput: &[]schemas.BifrostMessage{ + { + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ContentStr: stringPtr("Get documentation for React hooks using the get-library-docs tool with library ID '/facebook/react'.")}, + }, + }, + }, + Params: &schemas.ModelParameters{ + ToolChoice: &schemas.ToolChoice{ + ToolChoiceStr: stringPtr("auto"), + }, + }, + }) + + if bifrostErr != nil { + t.Fatalf("Chat completion request failed: %v", bifrostErr) + } + + if response == nil || len(response.Choices) == 0 { + t.Fatal("Expected response with choices") + } + + choice := response.Choices[0] + fmt.Printf("Response role: %s\n", choice.Message.Role) + + if choice.Message.AssistantMessage != nil && choice.Message.AssistantMessage.ToolCalls != nil { + fmt.Printf("Response contains %d tool calls\n", len(*choice.Message.AssistantMessage.ToolCalls)) + for i, toolCall := range *choice.Message.AssistantMessage.ToolCalls { + fmt.Printf(" Tool call %d: %s\n", i+1, *toolCall.Function.Name) + if *toolCall.Function.Name == "get-library-docs" { + t.Logf("Successfully called get-library-docs with arguments: %s", toolCall.Function.Arguments) + } + } + } else { + t.Log("No tool calls made - this might be expected if the model chooses not to call tools") + } + }) + + // Test 4: Use Context7 to get documentation and provide a final answer + t.Run("GetLibraryDocsAndAnswer_ThroughBifrost", func(t *testing.T) { + fmt.Println("=== CONTEXT7 GET DOCS AND ANSWER TEST ===") + + response, bifrostErr := client.ChatCompletionRequest(context.Background(), &schemas.BifrostRequest{ + Provider: schemas.OpenAI, + Model: "gpt-4o-mini", // A capable model is needed for this + Input: schemas.RequestInput{ + ChatCompletionInput: &[]schemas.BifrostMessage{ + { + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ContentStr: stringPtr("Can you explain what React's useReducer hook is and provide a simple code example? Use the available tools.")}, + }, + }, + }, + Params: &schemas.ModelParameters{ + ToolChoice: &schemas.ToolChoice{ + ToolChoiceStr: stringPtr("auto"), + }, + }, + }) + + if bifrostErr != nil { + t.Fatalf("Chat completion request failed: %v", bifrostErr) + } + + if response == nil || len(response.Choices) == 0 { + t.Fatal("Expected response with choices") + } + + choice := response.Choices[0] + fmt.Printf("Final response role: %s\n", choice.Message.Role) + + if choice.Message.Content.ContentStr == nil { + t.Fatal("Expected a final text response from the model, but got nil content") + } + + finalAnswer := *choice.Message.Content.ContentStr + fmt.Printf("Final answer from model:\n%s\n", finalAnswer) + + // Check for keywords that indicate the model successfully used the tool and synthesized an answer + expectedKeywords := []string{"usereducer", "hook", "state", "dispatch", "reducer"} + answerLower := strings.ToLower(finalAnswer) + + for _, keyword := range expectedKeywords { + if !strings.Contains(answerLower, keyword) { + t.Errorf("Final answer is missing expected keyword: '%s'", keyword) + } + } + + // Also check that no tool calls are in the *final* response, as they should have been handled + if choice.Message.AssistantMessage != nil && choice.Message.AssistantMessage.ToolCalls != nil { + t.Errorf("Expected final response to be a text message, but it contained tool calls") + } + }) + + // Cleanup + client.Cleanup() + fmt.Println("\n=== CONTEXT7 INTEGRATION TEST COMPLETED ===") +} + +func TestMCPPlugin_Maxim(t *testing.T) { + // Skip this test if no OPENAI_API_KEY is set + if os.Getenv("OPENAI_API_KEY") == "" { + t.Skip("Skipping Maxim MCP integration test: OPENAI_API_KEY not set") + } + + // Create MCP plugin with agentic mode enabled + plugin, err := NewMCPPlugin(MCPPluginConfig{ + AgenticMode: true, // Enable agentic flow + }, bifrost.NewDefaultLogger(schemas.LogLevelDebug)) + if err != nil { + t.Fatalf("Failed to create plugin: %v", err) + } + + // Connect to Maxim MCP server using uvx + err = plugin.ConnectToExternalMCP(ExternalMCPConfig{ + Name: "maxim-mcp", + ConnectionType: ConnectionTypeSTDIO, + StdioConfig: &StdioConfig{ + Command: "npx", + Args: []string{"-y", "@maximai/mcp-server@latest"}, + }, + }) + if err != nil { + t.Fatalf("Failed to connect to maxim-mcp MCP: %v", err) + } + + // Give some time for the external server to start and register tools + time.Sleep(3 * time.Second) + + // Initialize Bifrost with the MCP plugin + account := BaseAccount{} + client, err := bifrost.Init(schemas.BifrostConfig{ + Account: &account, + Plugins: []schemas.Plugin{plugin}, + Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug), + }) + if err != nil { + t.Fatalf("Failed to initialize Bifrost: %v", err) + } + + // Set the Bifrost client for agentic mode + plugin.SetBifrostClient(client) + + t.Run("SearchAndAnswer_ThroughBifrost", func(t *testing.T) { + fmt.Println("=== MAXIM MCP AND ANSWER TEST ===") + + response, bifrostErr := client.ChatCompletionRequest(context.Background(), &schemas.BifrostRequest{ + Provider: schemas.OpenAI, + Model: "gpt-4o-mini", // A capable model is needed for this + Input: schemas.RequestInput{ + ChatCompletionInput: &[]schemas.BifrostMessage{ + { + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ContentStr: stringPtr("fetch log-repository-entities from repository-id cma3pm2qn0668oeu5mdylurdv and workspace-id cm82w2kn1004r9wkqhsixdpzg")}, + }, + }, + }, + Params: &schemas.ModelParameters{ + ToolChoice: &schemas.ToolChoice{ + ToolChoiceStr: stringPtr("auto"), + }, + }, + }) + + if bifrostErr != nil { + t.Fatalf("Chat completion request failed: %v", bifrostErr) + } + + if response == nil || len(response.Choices) == 0 { + t.Fatal("Expected response with choices") + } + + // printChoicesAsJSON(t, response.Choices) + + choice := response.Choices[0] + fmt.Printf("Final response role: %s\n", choice.Message.Role) + + if choice.Message.Content.ContentStr == nil { + t.Fatal("Expected a final text response from the model, but got nil content") + } + + finalAnswer := *choice.Message.Content.ContentStr + fmt.Printf("Final answer from model:\n%s\n", finalAnswer) + + // Check for keywords that indicate the model successfully used the tool and synthesized an answer + // Using a broad check because the exact phrasing can vary. + answerLower := strings.ToLower(finalAnswer) + if !strings.Contains(answerLower, "bifrost") { + t.Errorf("Final answer is missing expected keyword: 'bifrost'") + } + + // Also check that no tool calls are in the *final* response, as they should have been handled + if choice.Message.AssistantMessage != nil && choice.Message.AssistantMessage.ToolCalls != nil && len(*choice.Message.AssistantMessage.ToolCalls) > 0 { + t.Errorf("Expected final response to be a text message, but it contained tool calls") + } + }) + + // Cleanup + client.Cleanup() + fmt.Println("\n=== MAXIM MCP TEST COMPLETED ===") +} + +func TestMCPPlugin_AgenticMode(t *testing.T) { + // Skip this test if no OPENAI_API_KEY is set + if os.Getenv("OPENAI_API_KEY") == "" { + t.Skip("Skipping agentic mode test: OPENAI_API_KEY not set") + } + + t.Run("NonAgenticMode", func(t *testing.T) { + // Test non-agentic mode (original behavior) + plugin, err := NewMCPPlugin(MCPPluginConfig{ + ServerPort: ":8383", + AgenticMode: false, // Disable agentic flow + }, bifrost.NewDefaultLogger(schemas.LogLevelDebug)) + if err != nil { + t.Fatalf("Failed to create MCP plugin: %v", err) + } + + // Register a simple test tool + testToolSchema := schemas.Tool{ + Type: "function", + Function: schemas.Function{ + Name: "test_tool", + Description: "A simple test tool", + Parameters: schemas.FunctionParameters{ + Type: "object", + Properties: map[string]interface{}{}, + Required: []string{}, + }, + }, + } + + err = RegisterTool(plugin, "test_tool", "A simple test tool", + func(args struct{}) (string, error) { + return "Test tool executed successfully", nil + }, testToolSchema, ToolExecutionPolicyAutoExecute) + if err != nil { + t.Fatalf("Failed to register test tool: %v", err) + } + + // Initialize Bifrost + account := BaseAccount{} + client, err := bifrost.Init(schemas.BifrostConfig{ + Account: &account, + Plugins: []schemas.Plugin{plugin}, + Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug), + }) + if err != nil { + t.Fatalf("Failed to initialize Bifrost: %v", err) + } + + // Test request that should trigger tool call + response, bifrostErr := client.ChatCompletionRequest(context.Background(), &schemas.BifrostRequest{ + Provider: schemas.OpenAI, + Model: "gpt-4o-mini", + Input: schemas.RequestInput{ + ChatCompletionInput: &[]schemas.BifrostMessage{ + { + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ContentStr: stringPtr("Use the test_tool")}, + }, + }, + }, + Params: &schemas.ModelParameters{ + ToolChoice: &schemas.ToolChoice{ + ToolChoiceStr: stringPtr("auto"), + }, + }, + }) + + if bifrostErr != nil { + t.Fatalf("Request failed: %v", bifrostErr) + } + + if response == nil || len(response.Choices) == 0 { + t.Fatal("Expected response with choices") + } + + // In non-agentic mode, the response should be a tool message + choice := response.Choices[0] + if choice.Message.Role != schemas.ModelChatMessageRoleTool { + t.Errorf("Expected tool message role, got: %s", choice.Message.Role) + } + + client.Cleanup() + }) + + t.Run("AgenticMode", func(t *testing.T) { + // Test agentic mode (new behavior) + plugin, err := NewMCPPlugin(MCPPluginConfig{ + ServerPort: ":8384", + AgenticMode: true, // Enable agentic flow + }, bifrost.NewDefaultLogger(schemas.LogLevelDebug)) + if err != nil { + t.Fatalf("Failed to create MCP plugin: %v", err) + } + + // Register a simple test tool + testToolSchema := schemas.Tool{ + Type: "function", + Function: schemas.Function{ + Name: "test_tool_agentic", + Description: "A simple test tool for agentic mode", + Parameters: schemas.FunctionParameters{ + Type: "object", + Properties: map[string]interface{}{}, + Required: []string{}, + }, + }, + } + + err = RegisterTool(plugin, "test_tool_agentic", "A simple test tool for agentic mode", + func(args struct{}) (string, error) { + return "Tool found the answer: 42", nil + }, testToolSchema, ToolExecutionPolicyAutoExecute) + if err != nil { + t.Fatalf("Failed to register test tool: %v", err) + } + + // Initialize Bifrost + account := BaseAccount{} + client, err := bifrost.Init(schemas.BifrostConfig{ + Account: &account, + Plugins: []schemas.Plugin{plugin}, + Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug), + }) + if err != nil { + t.Fatalf("Failed to initialize Bifrost: %v", err) + } + + // Set the Bifrost client for agentic mode + plugin.SetBifrostClient(client) + + // Test request that should trigger tool call and then synthesis + response, bifrostErr := client.ChatCompletionRequest(context.Background(), &schemas.BifrostRequest{ + Provider: schemas.OpenAI, + Model: "gpt-4o-mini", + Input: schemas.RequestInput{ + ChatCompletionInput: &[]schemas.BifrostMessage{ + { + Role: schemas.ModelChatMessageRoleUser, + Content: schemas.MessageContent{ContentStr: stringPtr("Use the test_tool_agentic to find the meaning of life")}, + }, + }, + }, + Params: &schemas.ModelParameters{ + ToolChoice: &schemas.ToolChoice{ + ToolChoiceStr: stringPtr("auto"), + }, + }, + }) + + if bifrostErr != nil { + t.Fatalf("Request failed: %v", bifrostErr) + } + + if response == nil || len(response.Choices) == 0 { + t.Fatal("Expected response with choices") + } + + // In agentic mode, the response should be an assistant message with synthesized content + choice := response.Choices[0] + if choice.Message.Role != schemas.ModelChatMessageRoleAssistant { + t.Errorf("Expected assistant message role, got: %s", choice.Message.Role) + } + + // Check that the response contains some indication of synthesis + if choice.Message.Content.ContentStr == nil { + t.Fatal("Expected synthesized text response") + } + + finalAnswer := *choice.Message.Content.ContentStr + fmt.Printf("Agentic synthesized response: %s\n", finalAnswer) + + // The response should ideally contain information from the tool result + answerLower := strings.ToLower(finalAnswer) + if !strings.Contains(answerLower, "42") { + t.Log("Warning: Synthesized response may not contain tool result data") + } + + client.Cleanup() + + }) + + fmt.Println("\n=== AGENTIC MODE TESTS COMPLETED ===") +} diff --git a/plugins/mcp/utils.go b/plugins/mcp/utils.go new file mode 100644 index 0000000000..94054b2564 --- /dev/null +++ b/plugins/mcp/utils.go @@ -0,0 +1,204 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "slices" + "strings" + + "github.com/maximhq/bifrost/core/schemas" + mcp_golang "github.com/metoro-io/mcp-golang" +) + +// Helper method to find MCP client for a tool (fixes repeated code and race condition) +func (p *MCPPlugin) findMCPClientForTool(toolName string) *PluginClient { + p.mu.RLock() + defer p.mu.RUnlock() + + for _, client := range p.clientMap { + if client.ToolMap != nil { + if _, isMCPTool := client.ToolMap[toolName]; isMCPTool { + return client + } + } + } + return nil +} + +// Helper method to create tool response message (fixes repeated code) +func (p *MCPPlugin) createToolResponseMessage(toolCall schemas.ToolCall, responseText string) schemas.BifrostMessage { + return schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleTool, + Content: schemas.MessageContent{ContentStr: &responseText}, + ToolMessage: &schemas.ToolMessage{ + ToolCallID: toolCall.ID, + }, + } +} + +// Helper method to extract text from MCP response (fixes repeated code) +func (p *MCPPlugin) extractTextFromMCPResponse(toolResponse *mcp_golang.ToolResponse, toolName string) string { + if toolResponse == nil { + return fmt.Sprintf("MCP tool '%s' executed successfully", toolName) + } + + var responseTextBuilder strings.Builder + if len(toolResponse.Content) > 0 { + for _, contentBlock := range toolResponse.Content { + if contentBlock.TextContent != nil && contentBlock.TextContent.Text != "" { + responseTextBuilder.WriteString(contentBlock.TextContent.Text) + responseTextBuilder.WriteString("\n") + } + } + } + + if responseTextBuilder.Len() > 0 { + return strings.TrimSpace(responseTextBuilder.String()) + } + return fmt.Sprintf("MCP tool '%s' executed successfully", toolName) +} + +// Helper method to execute a single tool (fixes repeated code and context issues) +func (p *MCPPlugin) executeSingleTool(ctx context.Context, toolCall schemas.ToolCall) (schemas.BifrostMessage, error) { + if toolCall.Function.Name == nil { + return schemas.BifrostMessage{}, fmt.Errorf("tool call missing function name") + } + toolName := *toolCall.Function.Name + + // Parse tool arguments + var arguments map[string]interface{} + if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &arguments); err != nil { + return schemas.BifrostMessage{}, fmt.Errorf("failed to parse tool arguments: %v", err) + } + + // Call the tool via MCP client -> MCP server + toolResponse, callErr := p.callTool(ctx, toolName, arguments) + if callErr != nil { + return schemas.BifrostMessage{}, fmt.Errorf("MCP tool call failed: %v", callErr) + } + + // Extract text from MCP response + responseText := p.extractTextFromMCPResponse(toolResponse, toolName) + + // Create tool response message + return p.createToolResponseMessage(toolCall, responseText), nil +} + +// convertMCPToolToBifrostSchema converts an MCP tool to Bifrost schema format +func convertMCPToolToBifrostSchema(mcpTool *mcp_golang.ToolRetType) schemas.Tool { + // Convert MCP tool schema to Bifrost tool schema + properties := make(map[string]interface{}) + required := []string{} + + if mcpTool.InputSchema != nil { + if schemaMap, ok := mcpTool.InputSchema.(map[string]interface{}); ok { + if props, ok := schemaMap["properties"].(map[string]interface{}); ok { + properties = props + } + if req, ok := schemaMap["required"].([]interface{}); ok { + for _, r := range req { + if reqStr, ok := r.(string); ok { + required = append(required, reqStr) + } + } + } + } + } + + // If no properties are defined, create an empty properties object + // This is required by OpenAI's function calling schema + if properties == nil { + properties = make(map[string]interface{}) + } + + description := "" + if mcpTool.Description != nil { + description = *mcpTool.Description + } + + return schemas.Tool{ + Type: "function", + Function: schemas.Function{ + Name: mcpTool.Name, + Description: description, + Parameters: schemas.FunctionParameters{ + Type: "object", + Properties: properties, + Required: required, + }, + }, + } +} + +// shouldIncludeClient determines if a client should be included based on filtering rules +func (p *MCPPlugin) shouldIncludeClient(clientName string, includeClients, excludeClients []string) bool { + // If includeClients is specified, only include those clients (whitelist mode) + if len(includeClients) > 0 { + for _, includeName := range includeClients { + if clientName == includeName { + return true + } + } + return false // Not in include list + } + + // If excludeClients is specified, exclude those clients (blacklist mode) + if len(excludeClients) > 0 { + for _, excludeName := range excludeClients { + if clientName == excludeName { + return false + } + } + } + + // Default: include all clients + return true +} + +// checkAgenticModeAvailable checks if agentic mode is properly configured and logs errors +func (p *MCPPlugin) checkAgenticModeAvailable() bool { + if p.agenticMode && p.bifrostClient == nil { + p.logger.Warn(LogPrefix + " Agentic mode is enabled but Bifrost client is not set. Falling back to normal execution.") + p.logger.Info(LogPrefix + " Hint: Call plugin.SetBifrostClient(bifrostInstance) to enable agentic mode.") + return false + } + return p.agenticMode && p.bifrostClient != nil +} + +// getToolExecutionPolicy returns the execution policy for a given tool +func (p *MCPPlugin) getToolExecutionPolicy(toolName, clientName string) ToolExecutionPolicy { + p.mu.RLock() + defer p.mu.RUnlock() + + // Then check client-specific configuration (external MCP tools) + client, exists := p.clientMap[clientName] + if !exists { + // Default: require approval for unknown clients + return ToolExecutionPolicyRequireApproval + } + + // Check for tool-specific policy first + if toolPolicy, exists := client.ExecutionConfig.ToolPolicies[toolName]; exists { + return toolPolicy + } + + // Fall back to client default policy + return client.ExecutionConfig.DefaultPolicy +} + +// shouldRequireApproval returns true if the tool requires user approval +func (p *MCPPlugin) shouldRequireApproval(toolName, clientName string) bool { + return p.getToolExecutionPolicy(toolName, clientName) == ToolExecutionPolicyRequireApproval +} + +// shouldSkipTool returns true if the tool should be skipped for this client +func (p *MCPPlugin) shouldSkipTool(toolName, clientName string) bool { + // ConnectToExternalMCP function already has the mutex lock + client, exists := p.clientMap[clientName] + if !exists { + return false // Don't skip if client config doesn't exist + } + + return slices.Contains(client.ExecutionConfig.ToolsToSkip, toolName) +}