Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 36 additions & 36 deletions core/bifrost.go
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ type Bifrost struct {
oauth2Provider schemas.OAuth2Provider // OAuth provider instance
logger schemas.Logger // logger instance, default logger is used if not provided
tracer atomic.Value // tracer for distributed tracing (stores schemas.Tracer, NoOpTracer if not configured)
mcpManager *mcp.MCPManager // MCP integration manager (nil if MCP not configured)
McpManager mcp.MCPManagerInterface // MCP integration manager (nil if MCP not configured)
mcpInitOnce sync.Once // Ensures MCP manager is initialized only once
dropExcessRequests atomic.Bool // If true, in cases where the queue is full, requests will not wait for the queue to be empty and will be dropped instead.
keySelector schemas.KeySelector // Custom key selector function
Expand Down Expand Up @@ -286,7 +286,7 @@ func Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error) {
}
}
codeMode := starlark.NewStarlarkCodeMode(codeModeConfig)
bifrost.mcpManager = mcp.NewMCPManager(bifrostCtx, mcpConfig, bifrost.oauth2Provider, bifrost.logger, codeMode)
bifrost.McpManager = mcp.NewMCPManager(bifrostCtx, mcpConfig, bifrost.oauth2Provider, bifrost.logger, codeMode)
bifrost.logger.Info("MCP integration initialized successfully")
})
}
Expand Down Expand Up @@ -713,8 +713,8 @@ func (bifrost *Bifrost) ChatCompletionRequest(ctx *schemas.BifrostContext, req *
}

// Check if we should enter agent mode
if bifrost.mcpManager != nil {
return bifrost.mcpManager.CheckAndExecuteAgentForChatRequest(
if bifrost.McpManager != nil {
return bifrost.McpManager.CheckAndExecuteAgentForChatRequest(
ctx,
req,
response,
Expand Down Expand Up @@ -810,8 +810,8 @@ func (bifrost *Bifrost) ResponsesRequest(ctx *schemas.BifrostContext, req *schem
}

// Check if we should enter agent mode
if bifrost.mcpManager != nil {
return bifrost.mcpManager.CheckAndExecuteAgentForResponsesRequest(
if bifrost.McpManager != nil {
return bifrost.McpManager.CheckAndExecuteAgentForResponsesRequest(
ctx,
req,
response,
Expand Down Expand Up @@ -2661,11 +2661,11 @@ func (bifrost *Bifrost) getProviderMutex(providerKey schemas.ModelProvider) *syn
// return args.Message, nil
// }, toolSchema)
func (bifrost *Bifrost) RegisterMCPTool(name, description string, handler func(args any) (string, error), toolSchema schemas.ChatTool) error {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return fmt.Errorf("MCP is not configured in this Bifrost instance")
}

return bifrost.mcpManager.RegisterTool(name, description, handler, toolSchema)
return bifrost.McpManager.RegisterTool(name, description, handler, toolSchema)
}

// IMPORTANT: Running the MCP client management operations (GetMCPClients, AddMCPClient, RemoveMCPClient, EditMCPClientTools)
Expand All @@ -2679,11 +2679,11 @@ func (bifrost *Bifrost) RegisterMCPTool(name, description string, handler func(a
// - []schemas.MCPClient: List of all MCP clients
// - error: Any retrieval error
func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return nil, fmt.Errorf("MCP is not configured in this Bifrost instance")
}

clients := bifrost.mcpManager.GetClients()
clients := bifrost.McpManager.GetClients()
clientsInConfig := make([]schemas.MCPClient, 0, len(clients))

for _, client := range clients {
Expand Down Expand Up @@ -2721,10 +2721,10 @@ func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
// Returns:
// - []schemas.ChatTool: List of available tools
func (bifrost *Bifrost) GetAvailableMCPTools(ctx context.Context) []schemas.ChatTool {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return nil
}
return bifrost.mcpManager.GetAvailableTools(ctx)
return bifrost.McpManager.GetAvailableTools(ctx)
}

// AddMCPClient adds a new MCP client to the Bifrost instance.
Expand All @@ -2744,7 +2744,7 @@ func (bifrost *Bifrost) GetAvailableMCPTools(ctx context.Context) []schemas.Chat
// ConnectionString: &url,
// })
func (bifrost *Bifrost) AddMCPClient(config *schemas.MCPClientConfig) error {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
// Use sync.Once to ensure thread-safe initialization
bifrost.mcpInitOnce.Do(func() {
// Initialize with empty config - client will be added via AddClient below
Expand All @@ -2763,16 +2763,16 @@ func (bifrost *Bifrost) AddMCPClient(config *schemas.MCPClientConfig) error {
// Create Starlark CodeMode for code execution (with default config)
starlark.SetLogger(bifrost.logger)
codeMode := starlark.NewStarlarkCodeMode(nil)
bifrost.mcpManager = mcp.NewMCPManager(bifrost.ctx, mcpConfig, bifrost.oauth2Provider, bifrost.logger, codeMode)
bifrost.McpManager = mcp.NewMCPManager(bifrost.ctx, mcpConfig, bifrost.oauth2Provider, bifrost.logger, codeMode)
})
}

// Handle case where initialization succeeded elsewhere but manager is still nil
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return fmt.Errorf("MCP manager is not initialized")
}

return bifrost.mcpManager.AddClient(config)
return bifrost.McpManager.AddClient(config)
}

// RemoveMCPClient removes an MCP client from the Bifrost instance.
Expand All @@ -2791,20 +2791,20 @@ func (bifrost *Bifrost) AddMCPClient(config *schemas.MCPClientConfig) error {
// log.Fatalf("Failed to remove MCP client: %v", err)
// }
func (bifrost *Bifrost) RemoveMCPClient(id string) error {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return fmt.Errorf("MCP is not configured in this Bifrost instance")
}

return bifrost.mcpManager.RemoveClient(id)
return bifrost.McpManager.RemoveClient(id)
}

// SetMCPManager sets the MCP manager for this Bifrost instance.
// This is primarily used for testing purposes to inject a custom MCP manager.
// This allows injecting a custom MCP manager implementation (e.g., for enterprise features).
//
// Parameters:
// - manager: The MCP manager to set
func (bifrost *Bifrost) SetMCPManager(manager *mcp.MCPManager) {
bifrost.mcpManager = manager
// - manager: The MCP manager to set (must implement MCPManagerInterface)
func (bifrost *Bifrost) SetMCPManager(manager mcp.MCPManagerInterface) {
bifrost.McpManager = manager
}
Comment on lines 2801 to 2808
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

# First, check if there's information about a PR stack or related changes
git log --oneline -20

Repository: maximhq/bifrost

Length of output: 83


🏁 Script executed:

# Find the SetMCPManager method to see the actual implementation
fd core/bifrost.go -x cat -n {} | head -100

Repository: maximhq/bifrost

Length of output: 41


🏁 Script executed:

# Get more context around lines 2801-2808 in core/bifrost.go
wc -l core/bifrost.go

Repository: maximhq/bifrost

Length of output: 79


🏁 Script executed:

# Search for how McpManager is used in the codebase
rg "McpManager" --type go -A 2 -B 2 | head -100

Repository: maximhq/bifrost

Length of output: 4959


🏁 Script executed:

# Check the MCPManagerInterface definition
rg "MCPManagerInterface" --type go -A 5 -B 2 | head -150

Repository: maximhq/bifrost

Length of output: 2341


🏁 Script executed:

# Search for actual calls to SetMCPManager
rg "SetMCPManager" --type go -B 3 -A 3

Repository: maximhq/bifrost

Length of output: 50371


🏁 Script executed:

# Check if there are any typed-nil patterns or reflect usage in the codebase
rg "reflect\." --type go | head -30

Repository: maximhq/bifrost

Length of output: 2613


🏁 Script executed:

# Get the exact location and context of SetMCPManager in bifrost.go
sed -n '2795,2815p' core/bifrost.go

Repository: maximhq/bifrost

Length of output: 756


🏁 Script executed:

# Check the MCPManagerInterface method signatures to understand what methods are called
fd core/mcp/interface.go -x cat -n | head -200

Repository: maximhq/bifrost

Length of output: 41


🏁 Script executed:

# Search for nil checks patterns in the codebase related to McpManager
rg "McpManager.*!= nil|McpManager.*== nil" --type go -B 2 -A 5

Repository: maximhq/bifrost

Length of output: 6262


🏁 Script executed:

# Check if there are any existing patterns in Go standard library or common patterns for validating interface implementations
rg "reflect.ValueOf.*Kind.*Ptr.*IsNil" --type go || echo "No match found in codebase"

Repository: maximhq/bifrost

Length of output: 85


🏁 Script executed:

# Look at how other setters in the codebase handle similar interface validation
rg "func.*Set.*interface\{}" --type go -A 3 | head -50

Repository: maximhq/bifrost

Length of output: 708


🏁 Script executed:

# Check if there are any other MCPManager validation patterns
fd core/mcp/ -name "*.go" -x cat -n {} | grep -A 10 -B 5 "nil" | head -80

Repository: maximhq/bifrost

Length of output: 289


🏁 Script executed:

# Check the MCPManager implementation to see if it's likely a typed-nil could exist
fd "manager.go" core/mcp/ -x head -100 {}

Repository: maximhq/bifrost

Length of output: 7046


🏁 Script executed:

# Search for how MCPManager is instantiated
rg "NewMCPManager" --type go -B 3 -A 5

Repository: maximhq/bifrost

Length of output: 3504


🏁 Script executed:

# Check if there are any tests or patterns showing typed-nil issues
rg "typed.nil|typed nil" --type go

Repository: maximhq/bifrost

Length of output: 41


Consider a simpler safeguard for SetMCPManager against nil or invalid implementations.

While a nil pointer stored in an interface value can pass != nil checks in Go, a more practical concern is ensuring the manager is properly initialized before use. However, the reflect-based approach is unconventional. A simpler alternative would be to document this expectation or add a basic nil check that converts to nil explicitly:

func (bifrost *Bifrost) SetMCPManager(manager mcp.MCPManagerInterface) {
+	if manager == nil {
+		bifrost.McpManager = nil
+		return
+	}
	bifrost.McpManager = manager
}

This covers the common case without requiring reflection. If stricter typed-nil validation is needed, that's better handled through code review and documentation of what valid implementations must provide.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// SetMCPManager sets the MCP manager for this Bifrost instance.
// This is primarily used for testing purposes to inject a custom MCP manager.
// This allows injecting a custom MCP manager implementation (e.g., for enterprise features).
//
// Parameters:
// - manager: The MCP manager to set
func (bifrost *Bifrost) SetMCPManager(manager *mcp.MCPManager) {
bifrost.mcpManager = manager
// - manager: The MCP manager to set (must implement MCPManagerInterface)
func (bifrost *Bifrost) SetMCPManager(manager mcp.MCPManagerInterface) {
bifrost.McpManager = manager
}
// SetMCPManager sets the MCP manager for this Bifrost instance.
// This allows injecting a custom MCP manager implementation (e.g., for enterprise features).
//
// Parameters:
// - manager: The MCP manager to set (must implement MCPManagerInterface)
func (bifrost *Bifrost) SetMCPManager(manager mcp.MCPManagerInterface) {
if manager == nil {
bifrost.McpManager = nil
return
}
bifrost.McpManager = manager
}
🤖 Prompt for AI Agents
In `@core/bifrost.go` around lines 2801 - 2808, The SetMCPManager method currently
assigns whatever is passed into bifrost.McpManager; add a simple nil guard so
that if the incoming manager is nil (or an untyped nil interface) you explicitly
set bifrost.McpManager = nil, otherwise assign the provided manager; locate the
SetMCPManager function and the MCPManagerInterface type to implement this check
and update the assignment accordingly and optionally add a short comment
documenting that callers must pass a properly initialized MCPManagerInterface.


// UpdateMCPClient updates the MCP client.
Expand All @@ -2824,11 +2824,11 @@ func (bifrost *Bifrost) SetMCPManager(manager *mcp.MCPManager) {
// ToolsToExecute: []string{"tool1", "tool2"},
// })
func (bifrost *Bifrost) EditMCPClient(id string, updatedConfig *schemas.MCPClientConfig) error {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return fmt.Errorf("MCP is not configured in this Bifrost instance")
}

return bifrost.mcpManager.EditClient(id, updatedConfig)
return bifrost.McpManager.EditClient(id, updatedConfig)
}

// ReconnectMCPClient attempts to reconnect an MCP client if it is disconnected.
Expand All @@ -2839,21 +2839,21 @@ func (bifrost *Bifrost) EditMCPClient(id string, updatedConfig *schemas.MCPClien
// Returns:
// - error: Any reconnection error
func (bifrost *Bifrost) ReconnectMCPClient(id string) error {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return fmt.Errorf("MCP is not configured in this Bifrost instance")
}

return bifrost.mcpManager.ReconnectClient(id)
return bifrost.McpManager.ReconnectClient(id)
}

// UpdateToolManagerConfig updates the tool manager config for the MCP manager.
// This allows for hot-reloading of the tool manager config at runtime.
func (bifrost *Bifrost) UpdateToolManagerConfig(maxAgentDepth int, toolExecutionTimeoutInSeconds int, codeModeBindingLevel string) error {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return fmt.Errorf("MCP is not configured in this Bifrost instance")
}

bifrost.mcpManager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
bifrost.McpManager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
MaxAgentDepth: maxAgentDepth,
ToolExecutionTimeout: time.Duration(toolExecutionTimeoutInSeconds) * time.Second,
CodeModeBindingLevel: schemas.CodeModeBindingLevel(codeModeBindingLevel),
Expand Down Expand Up @@ -3445,8 +3445,8 @@ func (bifrost *Bifrost) tryRequest(ctx *schemas.BifrostContext, req *schemas.Bif
}

// Add MCP tools to request if MCP is configured and requested
if bifrost.mcpManager != nil {
req = bifrost.mcpManager.AddToolsToRequest(ctx, req)
if bifrost.McpManager != nil {
req = bifrost.McpManager.AddToolsToRequest(ctx, req)
}

tracer := bifrost.getTracer()
Expand Down Expand Up @@ -3635,8 +3635,8 @@ func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schem
}

// Add MCP tools to request if MCP is configured and requested
if req.RequestType != schemas.SpeechStreamRequest && req.RequestType != schemas.TranscriptionStreamRequest && bifrost.mcpManager != nil {
req = bifrost.mcpManager.AddToolsToRequest(ctx, req)
if req.RequestType != schemas.SpeechStreamRequest && req.RequestType != schemas.TranscriptionStreamRequest && bifrost.McpManager != nil {
req = bifrost.McpManager.AddToolsToRequest(ctx, req)
}

tracer := bifrost.getTracer()
Expand Down Expand Up @@ -4411,7 +4411,7 @@ func (bifrost *Bifrost) handleProviderStreamRequest(provider schemas.Provider, r
// - *schemas.BifrostMCPResponse: The MCP response after all hooks
// - *schemas.BifrostError: Any execution error
func (bifrost *Bifrost) handleMCPToolExecution(ctx *schemas.BifrostContext, mcpRequest *schemas.BifrostMCPRequest, requestType schemas.RequestType) (*schemas.BifrostMCPResponse, *schemas.BifrostError) {
if bifrost.mcpManager == nil {
if bifrost.McpManager == nil {
return nil, &schemas.BifrostError{
IsBifrostError: false,
Error: &schemas.ErrorField{
Expand Down Expand Up @@ -4470,7 +4470,7 @@ func (bifrost *Bifrost) handleMCPToolExecution(ctx *schemas.BifrostContext, mcpR
}

// Execute tool with modified request
result, err := bifrost.mcpManager.ExecuteToolCall(ctx, preReq)
result, err := bifrost.McpManager.ExecuteToolCall(ctx, preReq)

// Prepare MCP response and error for post-hooks
var mcpResp *schemas.BifrostMCPResponse
Expand Down Expand Up @@ -5278,8 +5278,8 @@ func (bifrost *Bifrost) Shutdown() {
})

// Cleanup MCP manager
if bifrost.mcpManager != nil {
err := bifrost.mcpManager.Cleanup()
if bifrost.McpManager != nil {
err := bifrost.McpManager.Cleanup()
if err != nil {
bifrost.logger.Warn("Error cleaning up MCP manager: %s", err.Error())
}
Expand Down
4 changes: 2 additions & 2 deletions core/mcp/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -182,8 +182,8 @@ func executeAgent(
toolName := *toolCall.Function.Name
client := clientManager.GetClientForTool(toolName)
if client == nil {
// Allow code mode list and read tool tools
if toolName == ToolTypeListToolFiles || toolName == ToolTypeReadToolFile {
// Allow code mode list, read, and docs tools (all read-only operations)
if toolName == ToolTypeListToolFiles || toolName == ToolTypeReadToolFile || toolName == ToolTypeGetToolDocs {
autoExecutableTools = append(autoExecutableTools, toolCall)
logger.Debug("Tool %s can be auto-executed", toolName)
continue
Expand Down
47 changes: 30 additions & 17 deletions core/mcp/codemode/starlark/readfile.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,15 +27,19 @@ func (s *StarlarkCodeMode) createReadToolFileTool() schemas.ChatTool {
"The function performs case-insensitive matching and removes the .pyi extension. " +
"Each tool can be accessed in code via: serverName.tool_name(param=value). " +
"If the compact signature is not enough to understand a tool, use getToolDocs for detailed documentation. " +
"Workflow: listToolFiles -> readToolFile -> (optional) getToolDocs -> executeToolCode."
"Workflow: listToolFiles -> readToolFile -> (optional) getToolDocs -> executeToolCode. " +
"IMPORTANT: If the response header shows 'Total lines: X (this is the complete file)', " +
"do NOT call this tool again with startLine/endLine - you already have the complete file."
} else {
fileNameDescription = "The virtual filename from listToolFiles in format: servers/<serverName>/<toolName>.pyi (e.g., 'calculator/add.pyi')"
toolDescription = "Reads a virtual .pyi stub file for a specific tool, returning its compact Python function signature. " +
"The fileName should be in format servers/<serverName>/<toolName>.pyi as listed by listToolFiles. " +
"The function performs case-insensitive matching and removes the .pyi extension. " +
"The tool can be accessed in code via: serverName.tool_name(param=value). " +
"If the compact signature is not enough to understand the tool, use getToolDocs for detailed documentation. " +
"Workflow: listToolFiles -> readToolFile -> (optional) getToolDocs -> executeToolCode."
"Workflow: listToolFiles -> readToolFile -> (optional) getToolDocs -> executeToolCode. " +
"IMPORTANT: If the response header shows 'Total lines: X (this is the complete file)', " +
"do NOT call this tool again with startLine/endLine - you already have the complete file."
}

readToolFileProps := schemas.OrderedMap{
Expand All @@ -45,11 +49,11 @@ func (s *StarlarkCodeMode) createReadToolFileTool() schemas.ChatTool {
},
"startLine": map[string]interface{}{
"type": "number",
"description": "Optional 1-based starting line number for partial file read (inclusive). Note: Line numbers start at 1, not 0. The first line is line 1.",
"description": "Optional 1-based starting line number for partial file read. Usually not needed - omit to read the entire file. Files are typically small (under 50 lines).",
},
"endLine": map[string]interface{}{
"type": "number",
"description": "Optional 1-based ending line number for partial file read (inclusive)",
"description": "Optional 1-based ending line number for partial file read. Usually not needed - omit to read the entire file. Will be clamped to actual file size if too large.",
},
}
return schemas.ChatTool{
Expand Down Expand Up @@ -197,6 +201,12 @@ func (s *StarlarkCodeMode) handleReadToolFile(ctx context.Context, toolCall sche
lines := strings.Split(fileContent, "\n")
totalLines := len(lines)

// Prepend total lines info so LLM knows the file size upfront
fileContent = fmt.Sprintf("# Total lines: %d (this is the complete file, no need to paginate)\n%s", totalLines+1, fileContent)
// Recalculate lines after prepending
lines = strings.Split(fileContent, "\n")
totalLines = len(lines)

// Handle line slicing if provided
var startLine, endLine *int
if sl, ok := arguments["startLine"].(float64); ok {
Expand All @@ -218,21 +228,23 @@ func (s *StarlarkCodeMode) handleReadToolFile(ctx context.Context, toolCall sche
end = *endLine
}

// Validate line numbers
if start < 1 || start > totalLines {
errorMsg := fmt.Sprintf("Invalid startLine: %d. Must be between 1 and %d (total lines in file). Provided: startLine=%d, endLine=%v, totalLines=%d",
start, totalLines, start, endLine, totalLines)
return createToolResponseMessage(toolCall, errorMsg), nil
// Clamp values to valid range instead of erroring
// This handles cases where LLM requests more lines than exist
if start < 1 {
start = 1
}
if start > totalLines {
start = totalLines
}
if end < 1 {
end = 1
}
if end < 1 || end > totalLines {
errorMsg := fmt.Sprintf("Invalid endLine: %d. Must be between 1 and %d (total lines in file). Provided: startLine=%d, endLine=%d, totalLines=%d",
end, totalLines, start, end, totalLines)
return createToolResponseMessage(toolCall, errorMsg), nil
if end > totalLines {
end = totalLines
}
if start > end {
errorMsg := fmt.Sprintf("Invalid line range: startLine (%d) must be less than or equal to endLine (%d). Total lines in file: %d",
start, end, totalLines)
return createToolResponseMessage(toolCall, errorMsg), nil
// If start > end after clamping, just return the start line
end = start
}

// Slice lines (convert to 0-based indexing)
Expand Down Expand Up @@ -289,7 +301,8 @@ func generateCompactSignatures(clientName string, tools []schemas.ChatTool, isTo
sb.WriteString(fmt.Sprintf("# %s server tools\n", clientName))
}
sb.WriteString(fmt.Sprintf("# Usage: %s.tool_name(param=value)\n", clientName))
sb.WriteString(fmt.Sprintf("# For detailed docs: use getToolDocs(server=\"%s\", tool=\"tool_name\")\n\n", clientName))
sb.WriteString(fmt.Sprintf("# For detailed docs: use getToolDocs(server=\"%s\", tool=\"tool_name\")\n", clientName))
sb.WriteString("# Note: Descriptions may be truncated. Use getToolDocs for full details.\n\n")

for _, tool := range tools {
if tool.Function == nil || tool.Function.Name == "" {
Expand Down
Loading