Skip to content
Open

v1.4.0 #1153

Show file tree
Hide file tree
Changes from 10 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
4 changes: 4 additions & 0 deletions .editorconfig
Original file line number Diff line number Diff line change
Expand Up @@ -5,5 +5,9 @@ insert_final_newline = false
end_of_line = lf
charset = utf-8

[*.go]
indent_style = tab
indent_size = 4

[*.{js,jsx,ts,tsx,mjs,json,md,css,scss,html}]
insert_final_newline = false
138 changes: 103 additions & 35 deletions core/bifrost.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"time"

"github.com/google/uuid"
"github.com/maximhq/bifrost/core/mcp"
"github.com/maximhq/bifrost/core/providers/anthropic"
"github.com/maximhq/bifrost/core/providers/azure"
"github.com/maximhq/bifrost/core/providers/bedrock"
Expand Down Expand Up @@ -66,7 +67,8 @@ type Bifrost struct {
pluginPipelinePool sync.Pool // Pool for PluginPipeline objects
bifrostRequestPool sync.Pool // Pool for BifrostRequest objects
logger schemas.Logger // logger instance, default logger is used if not provided
mcpManager *MCPManager // MCP integration manager (nil if MCP not configured)
mcpManager *mcp.MCPManager // 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 @@ -179,13 +181,10 @@ func Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error) {

// Initialize MCP manager if configured
if config.MCPConfig != nil {
mcpManager, err := newMCPManager(bifrostCtx, *config.MCPConfig, bifrost.logger)
if err != nil {
bifrost.logger.Warn(fmt.Sprintf("failed to initialize MCP manager: %v", err))
} else {
bifrost.mcpManager = mcpManager
bifrost.mcpInitOnce.Do(func() {
bifrost.mcpManager = mcp.NewMCPManager(bifrostCtx, *config.MCPConfig, bifrost.logger)
bifrost.logger.Info("MCP integration initialized successfully")
}
})
}

// Create buffered channels for each provider and start workers
Expand Down Expand Up @@ -539,8 +538,7 @@ func (bifrost *Bifrost) TextCompletionStreamRequest(ctx context.Context, req *sc
return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// ChatCompletionRequest sends a chat completion request to the specified provider.
func (bifrost *Bifrost) ChatCompletionRequest(ctx context.Context, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
func (bifrost *Bifrost) makeChatCompletionRequest(ctx context.Context, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
if req == nil {
return nil, &schemas.BifrostError{
IsBifrostError: false,
Expand Down Expand Up @@ -574,10 +572,35 @@ func (bifrost *Bifrost) ChatCompletionRequest(ctx context.Context, req *schemas.
if err != nil {
return nil, err
}
//TODO: Release the response

return response.ChatResponse, nil
}

// ChatCompletionRequest sends a chat completion request to the specified provider.
func (bifrost *Bifrost) ChatCompletionRequest(ctx context.Context, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
// If ctx is nil, use the bifrost context (defensive check for mcp agent mode)
if ctx == nil {
ctx = bifrost.ctx
}

response, err := bifrost.makeChatCompletionRequest(ctx, req)
if err != nil {
return nil, err
}

// Check if we should enter agent mode
if bifrost.mcpManager != nil {
return bifrost.mcpManager.CheckAndExecuteAgentForChatRequest(
&ctx,
req,
response,
bifrost.makeChatCompletionRequest,
)
}
Comment on lines +592 to +599
Copy link
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

# Examine the code at lines 592-599 in core/bifrost.go
sed -n '585,605p' core/bifrost.go

Repository: maximhq/bifrost

Length of output: 655


🏁 Script executed:

# Find the CheckAndExecuteAgentForChatRequest function definition
rg -A 15 "func.*CheckAndExecuteAgentForChatRequest" --type go

Repository: maximhq/bifrost

Length of output: 925


🏁 Script executed:

# Check the full context of the bifrost.mcpManager calls
sed -n '585,605p' core/bifrost.go
sed -n '682,702p' core/bifrost.go

Repository: maximhq/bifrost

Length of output: 1252


🏁 Script executed:

# Check how the ctx pointer parameter is used inside CheckAndExecuteAgentForChatRequest
rg -A 50 "func.*CheckAndExecuteAgentForChatRequest" core/mcp/mcp.go | head -80

Repository: maximhq/bifrost

Length of output: 2423


🏁 Script executed:

# Look for context dereferences (*ctx) in the mcp.go file
rg "\*ctx" core/mcp/mcp.go | head -20

Repository: maximhq/bifrost

Length of output: 41


🏁 Script executed:

# Check if context is modified (WithValue, WithCancel, etc.) in this function
rg "context\.With" core/mcp/mcp.go

Repository: maximhq/bifrost

Length of output: 41


🏁 Script executed:

# Check the ExecuteAgentForChatRequest signature to see what parameter type it expects
rg -B 2 -A 15 "func.*ExecuteAgentForChatRequest" core/mcp/ --type go

Repository: maximhq/bifrost

Length of output: 5543


🏁 Script executed:

# Check if the toolsManager methods dereference or modify context
rg "\*ctx|context\.With" core/mcp/tools_manager.go 2>/dev/null || echo "File not found, searching recursively"
rg "\*ctx|context\.With" core/mcp/ --type go

Repository: maximhq/bifrost

Length of output: 1497


Remove pointer-to-context pattern; it is an anti-pattern that does not actually propagate updates back to the caller.

The code passes &ctx to CheckAndExecuteAgentForChatRequest, which dereferences it as *ctx = context.WithValue(...) in core/mcp/agent.go. However, context.WithValue() returns a new context instance, not a modification of the existing one. Reassigning via the pointer only affects the dereferenced value within the function's call chain—it does not propagate changes back to the original caller's context variable. This pattern creates a false impression of context sharing and violates Go's immutable context semantics.

Refactor to use standard context patterns: pass ctx by value and return the modified context, or use a context that is already initialized with the necessary values before being passed.

🤖 Prompt for AI Agents
In core/bifrost.go around lines 592 to 599, the code is passing &ctx (a
pointer-to-context) into CheckAndExecuteAgentForChatRequest which is an
anti-pattern because context.WithValue returns a new context and changes won't
propagate back via pointer; instead change the call to pass ctx by value and
update CheckAndExecuteAgentForChatRequest to return the modified context (or
require the context to be pre-initialized) so the caller can capture the
returned context, remove the pointer parameter from the function signature and
all call sites, and update the implementation to accept and return
context.Context (or modify callers to set up context values before calling).


return response, nil
}

// ChatCompletionStreamRequest sends a chat completion stream request to the specified provider.
func (bifrost *Bifrost) ChatCompletionStreamRequest(ctx context.Context, req *schemas.BifrostChatRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
if req == nil {
Expand Down Expand Up @@ -612,8 +635,7 @@ func (bifrost *Bifrost) ChatCompletionStreamRequest(ctx context.Context, req *sc
return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// ResponsesRequest sends a responses request to the specified provider.
func (bifrost *Bifrost) ResponsesRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
func (bifrost *Bifrost) makeResponsesRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
if req == nil {
return nil, &schemas.BifrostError{
IsBifrostError: false,
Expand Down Expand Up @@ -647,10 +669,34 @@ func (bifrost *Bifrost) ResponsesRequest(ctx context.Context, req *schemas.Bifro
if err != nil {
return nil, err
}
//TODO: Release the response
return response.ResponsesResponse, nil
}

// ResponsesRequest sends a responses request to the specified provider.
func (bifrost *Bifrost) ResponsesRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
// If ctx is nil, use the bifrost context (defensive check for mcp agent mode)
if ctx == nil {
ctx = bifrost.ctx
}

response, err := bifrost.makeResponsesRequest(ctx, req)
if err != nil {
return nil, err
}

// Check if we should enter agent mode
if bifrost.mcpManager != nil {
return bifrost.mcpManager.CheckAndExecuteAgentForResponsesRequest(
&ctx,
req,
response,
bifrost.makeResponsesRequest,
)
}

return response, nil
}

// ResponsesStreamRequest sends a responses stream request to the specified provider.
func (bifrost *Bifrost) ResponsesStreamRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
if req == nil {
Expand Down Expand Up @@ -1644,7 +1690,7 @@ func (bifrost *Bifrost) RegisterMCPTool(name, description string, handler func(a
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)
}

// ExecuteMCPTool executes an MCP tool call and returns the result as a tool message.
Expand All @@ -1670,13 +1716,12 @@ func (bifrost *Bifrost) ExecuteMCPTool(ctx context.Context, toolCall schemas.Cha
}
}

result, err := bifrost.mcpManager.executeTool(ctx, toolCall)
result, err := bifrost.mcpManager.ExecuteTool(ctx, toolCall)
if err != nil {
return nil, &schemas.BifrostError{
IsBifrostError: false,
Error: &schemas.ErrorField{
Message: err.Error(),
Error: err,
},
ExtraFields: schemas.BifrostErrorExtraFields{
RequestType: schemas.ChatCompletionRequest, // MCP tools are used with chat completions
Expand All @@ -1702,12 +1747,9 @@ func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
return nil, fmt.Errorf("MCP is not configured in this Bifrost instance")
}

clients, err := bifrost.mcpManager.GetClients()
if err != nil {
return nil, err
}

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

for _, client := range clients {
tools := make([]schemas.ChatToolFunction, 0, len(client.ToolMap))
for _, tool := range client.ToolMap {
Expand Down Expand Up @@ -1735,6 +1777,17 @@ func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
return clientsInConfig, nil
}

// GetAvailableTools returns the available tools for the given context.
//
// Returns:
// - []schemas.ChatTool: List of available tools
func (bifrost *Bifrost) GetAvailableMCPTools(ctx context.Context) []schemas.ChatTool {
if bifrost.mcpManager == nil {
return nil
}
return bifrost.mcpManager.GetAvailableTools(ctx)
}

// AddMCPClient adds a new MCP client to the Bifrost instance.
// This allows for dynamic MCP client management at runtime.
//
Expand All @@ -1753,13 +1806,17 @@ func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
// })
func (bifrost *Bifrost) AddMCPClient(config schemas.MCPClientConfig) error {
if bifrost.mcpManager == nil {
manager := &MCPManager{
ctx: bifrost.ctx,
clientMap: make(map[string]*MCPClient),
logger: bifrost.logger,
}
// Use sync.Once to ensure thread-safe initialization
bifrost.mcpInitOnce.Do(func() {
bifrost.mcpManager = mcp.NewMCPManager(bifrost.ctx, schemas.MCPConfig{
ClientConfigs: []schemas.MCPClientConfig{config},
}, bifrost.logger)
})
}

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

return bifrost.mcpManager.AddClient(config)
Expand Down Expand Up @@ -1827,6 +1884,20 @@ func (bifrost *Bifrost) ReconnectMCPClient(id string) error {
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) error {
if bifrost.mcpManager == nil {
return fmt.Errorf("MCP is not configured in this Bifrost instance")
}

bifrost.mcpManager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
MaxAgentDepth: maxAgentDepth,
ToolExecutionTimeout: time.Duration(toolExecutionTimeoutInSeconds) * time.Second,
})
return nil
}

// PROVIDER MANAGEMENT

// createBaseProvider creates a provider based on the base provider type
Expand Down Expand Up @@ -2332,11 +2403,8 @@ func (bifrost *Bifrost) tryRequest(ctx context.Context, req *schemas.BifrostRequ
}

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

pipeline := bifrost.getPluginPipeline()
Expand Down Expand Up @@ -2452,7 +2520,7 @@ func (bifrost *Bifrost) tryStreamRequest(ctx context.Context, req *schemas.Bifro

// 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.addMCPToolsToBifrostRequest(ctx, req)
req = bifrost.mcpManager.AddToolsToRequest(ctx, req)
}

pipeline := bifrost.getPluginPipeline()
Expand Down Expand Up @@ -2644,7 +2712,7 @@ func executeRequestWithRetries[T any](

// Calculate and apply backoff
backoff := calculateBackoff(attempts-1, config)
logger.Debug("sleeping for %s", backoff)
logger.Debug("sleeping for %s before retry", backoff)
time.Sleep(backoff)
}

Expand Down Expand Up @@ -3463,7 +3531,7 @@ func (bifrost *Bifrost) Shutdown() {

// Cleanup MCP manager
if bifrost.mcpManager != nil {
err := bifrost.mcpManager.cleanup()
err := bifrost.mcpManager.Cleanup()
if err != nil {
bifrost.logger.Warn(fmt.Sprintf("Error cleaning up MCP manager: %s", err.Error()))
}
Expand Down
34 changes: 20 additions & 14 deletions core/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -5,23 +5,26 @@ go 1.25.5
require (
github.com/aws/aws-sdk-go-v2 v1.41.0
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.4
github.com/aws/aws-sdk-go-v2/config v1.32.5
github.com/aws/aws-sdk-go-v2/credentials v1.19.5
github.com/aws/aws-sdk-go-v2/service/s3 v1.93.2
github.com/aws/aws-sdk-go-v2/config v1.32.6
github.com/aws/aws-sdk-go-v2/credentials v1.19.6
github.com/aws/aws-sdk-go-v2/service/s3 v1.94.0
github.com/aws/smithy-go v1.24.0
github.com/bytedance/sonic v1.14.1
github.com/bytedance/sonic v1.14.2
github.com/clarkmcc/go-typescript v0.7.0
github.com/dop251/goja v0.0.0-20251103141225-af2ceb9156d7
github.com/google/uuid v1.6.0
github.com/hajimehoshi/go-mp3 v0.3.4
github.com/mark3labs/mcp-go v0.41.1
github.com/mark3labs/mcp-go v0.43.2
github.com/rs/zerolog v1.34.0
github.com/stretchr/testify v1.11.1
github.com/valyala/fasthttp v1.67.0
golang.org/x/oauth2 v0.32.0
golang.org/x/text v0.31.0
github.com/valyala/fasthttp v1.68.0
golang.org/x/oauth2 v0.34.0
golang.org/x/text v0.32.0
)

require (
cloud.google.com/go/compute/metadata v0.9.0 // indirect
github.com/Masterminds/semver/v3 v3.3.1 // indirect
github.com/andybalholm/brotli v1.2.0 // indirect
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.16 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.16 // indirect
Expand All @@ -33,17 +36,20 @@ require (
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.16 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.16 // indirect
github.com/aws/aws-sdk-go-v2/service/signin v1.0.4 // indirect
github.com/aws/aws-sdk-go-v2/service/sso v1.30.7 // indirect
github.com/aws/aws-sdk-go-v2/service/sso v1.30.8 // indirect
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.12 // indirect
github.com/aws/aws-sdk-go-v2/service/sts v1.41.5 // indirect
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/buger/jsonparser v1.1.1 // indirect
github.com/bytedance/gopkg v0.1.3 // indirect
github.com/bytedance/sonic/loader v0.3.0 // indirect
github.com/bytedance/sonic/loader v0.4.0 // indirect
github.com/cloudwego/base64x v0.1.6 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/dlclark/regexp2 v1.11.4 // indirect
github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect
github.com/google/pprof v0.0.0-20241029153458-d1b30febd7db // indirect
github.com/invopop/jsonschema v0.13.0 // indirect
github.com/klauspost/compress v1.18.1 // indirect
github.com/klauspost/compress v1.18.2 // indirect
github.com/klauspost/cpuid/v2 v2.3.0 // indirect
github.com/mailru/easyjson v0.9.1 // indirect
github.com/mattn/go-colorable v0.1.14 // indirect
Expand All @@ -55,9 +61,9 @@ require (
github.com/valyala/bytebufferpool v1.0.0 // indirect
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
golang.org/x/arch v0.22.0 // indirect
golang.org/x/net v0.47.0 // indirect
golang.org/x/sys v0.38.0 // indirect
golang.org/x/arch v0.23.0 // indirect
golang.org/x/net v0.48.0 // indirect
golang.org/x/sys v0.39.0 // indirect
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
Loading
Loading