Skip to content

Commit 57671fd

Browse files
feat: add option to disable automatic MCP tool injection per request
1 parent ea84c03 commit 57671fd

File tree

22 files changed

+171
-27
lines changed

22 files changed

+171
-27
lines changed

core/bifrost.go

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3174,7 +3174,7 @@ func (bifrost *Bifrost) getProviderMutex(providerKey schemas.ModelProvider) *syn
31743174
// }, toolSchema)
31753175
func (bifrost *Bifrost) RegisterMCPTool(name, description string, handler func(args any) (string, error), toolSchema schemas.ChatTool) error {
31763176
if bifrost.MCPManager == nil {
3177-
return fmt.Errorf("MCP is not configured in this Bifrost instance")
3177+
return fmt.Errorf("mcp is not configured in this bifrost instance")
31783178
}
31793179

31803180
return bifrost.MCPManager.RegisterTool(name, description, handler, toolSchema)
@@ -3192,7 +3192,7 @@ func (bifrost *Bifrost) RegisterMCPTool(name, description string, handler func(a
31923192
// - error: Any retrieval error
31933193
func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
31943194
if bifrost.MCPManager == nil {
3195-
return nil, fmt.Errorf("MCP is not configured in this Bifrost instance")
3195+
return nil, fmt.Errorf("mcp is not configured in this bifrost instance")
31963196
}
31973197

31983198
clients := bifrost.MCPManager.GetClients()
@@ -3303,7 +3303,7 @@ func (bifrost *Bifrost) AddMCPClient(config *schemas.MCPClientConfig) error {
33033303
// }
33043304
func (bifrost *Bifrost) RemoveMCPClient(id string) error {
33053305
if bifrost.MCPManager == nil {
3306-
return fmt.Errorf("MCP is not configured in this Bifrost instance")
3306+
return fmt.Errorf("mcp is not configured in this bifrost instance")
33073307
}
33083308

33093309
return bifrost.MCPManager.RemoveClient(id)
@@ -3336,7 +3336,7 @@ func (bifrost *Bifrost) SetMCPManager(manager mcp.MCPManagerInterface) {
33363336
// })
33373337
func (bifrost *Bifrost) UpdateMCPClient(id string, updatedConfig *schemas.MCPClientConfig) error {
33383338
if bifrost.MCPManager == nil {
3339-
return fmt.Errorf("MCP is not configured in this Bifrost instance")
3339+
return fmt.Errorf("mcp is not configured in this bifrost instance")
33403340
}
33413341

33423342
return bifrost.MCPManager.UpdateClient(id, updatedConfig)
@@ -3351,7 +3351,7 @@ func (bifrost *Bifrost) UpdateMCPClient(id string, updatedConfig *schemas.MCPCli
33513351
// - error: Any reconnection error
33523352
func (bifrost *Bifrost) ReconnectMCPClient(id string) error {
33533353
if bifrost.MCPManager == nil {
3354-
return fmt.Errorf("MCP is not configured in this Bifrost instance")
3354+
return fmt.Errorf("mcp is not configured in this bifrost instance")
33553355
}
33563356

33573357
return bifrost.MCPManager.ReconnectClient(id)
@@ -3361,7 +3361,7 @@ func (bifrost *Bifrost) ReconnectMCPClient(id string) error {
33613361
// This allows for hot-reloading of the tool manager config at runtime.
33623362
func (bifrost *Bifrost) UpdateToolManagerConfig(maxAgentDepth int, toolExecutionTimeoutInSeconds int, codeModeBindingLevel string) error {
33633363
if bifrost.MCPManager == nil {
3364-
return fmt.Errorf("MCP is not configured in this Bifrost instance")
3364+
return fmt.Errorf("mcp is not configured in this bifrost instance")
33653365
}
33663366

33673367
bifrost.MCPManager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
@@ -3372,6 +3372,16 @@ func (bifrost *Bifrost) UpdateToolManagerConfig(maxAgentDepth int, toolExecution
33723372
return nil
33733373
}
33743374

3375+
// UpdateMCPDisableAutoToolInject updates the auto tool injection toggle at runtime.
3376+
// This allows for hot-reloading without a server restart.
3377+
func (bifrost *Bifrost) UpdateMCPDisableAutoToolInject(val bool) error {
3378+
if bifrost.MCPManager == nil {
3379+
return fmt.Errorf("mcp is not configured in this bifrost instance")
3380+
}
3381+
bifrost.MCPManager.UpdateDisableAutoToolInject(val)
3382+
return nil
3383+
}
3384+
33753385
// PROVIDER MANAGEMENT
33763386

33773387
// createBaseProvider creates a provider based on the base provider type
@@ -5086,7 +5096,7 @@ func (bifrost *Bifrost) handleMCPToolExecution(ctx *schemas.BifrostContext, mcpR
50865096
return nil, &schemas.BifrostError{
50875097
IsBifrostError: false,
50885098
Error: &schemas.ErrorField{
5089-
Message: "MCP is not configured in this Bifrost instance",
5099+
Message: "mcp is not configured in this bifrost instance",
50905100
},
50915101
ExtraFields: schemas.BifrostErrorExtraFields{
50925102
RequestType: requestType,

core/changelog.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
1+
- feat: add DisableAutoToolInject to MCPToolManagerConfig to suppress automatic MCP tool injection per request
12
- feat: add Filename field to TranscriptionInput schema to carry original filename through the request pipeline
23
- fix: add AudioFilenameFromBytes utility to detect audio format from file headers with mp3 fallback

core/mcp/interface.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@ type MCPManagerInterface interface {
2525
// UpdateToolManagerConfig updates the configuration for the tool manager
2626
UpdateToolManagerConfig(config *schemas.MCPToolManagerConfig)
2727

28+
// UpdateDisableAutoToolInject updates the global auto tool injection toggle at runtime
29+
UpdateDisableAutoToolInject(val bool)
30+
2831
// Agent Mode Operations
2932
// CheckAndExecuteAgentForChatRequest handles agent mode for Chat Completions API
3033
CheckAndExecuteAgentForChatRequest(

core/mcp/mcp.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,13 @@ func (m *MCPManager) UpdateToolManagerConfig(config *schemas.MCPToolManagerConfi
174174
m.toolsManager.UpdateConfig(config)
175175
}
176176

177+
// UpdateDisableAutoToolInject updates the global auto tool injection toggle at runtime.
178+
// When val is true, MCP tools are not injected into requests unless an explicit
179+
// context filter (MCPContextKeyIncludeTools or MCPContextKeyIncludeClients) is set.
180+
func (m *MCPManager) UpdateDisableAutoToolInject(val bool) {
181+
m.toolsManager.UpdateDisableAutoToolInject(val)
182+
}
183+
177184
// CheckAndExecuteAgentForChatRequest checks if the chat response contains tool calls,
178185
// and if so, executes agent mode to handle the tool calls iteratively. If no tool calls
179186
// are present, it returns the original response unchanged.

core/mcp/toolmanager.go

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,10 @@ type PluginPipeline interface {
3030

3131
// ToolsManager manages MCP tool execution and agent mode.
3232
type ToolsManager struct {
33-
toolExecutionTimeout atomic.Value
34-
maxAgentDepth atomic.Int32
35-
clientManager ClientManager
33+
toolExecutionTimeout atomic.Value
34+
maxAgentDepth atomic.Int32
35+
disableAutoToolInject atomic.Bool
36+
clientManager ClientManager
3637
logger schemas.Logger
3738
agentModeExecutor *AgentModeExecutor
3839

@@ -146,6 +147,7 @@ func NewToolsManagerWithCodeMode(
146147
// Initialize atomic values
147148
manager.toolExecutionTimeout.Store(config.ToolExecutionTimeout)
148149
manager.maxAgentDepth.Store(int32(config.MaxAgentDepth))
150+
manager.disableAutoToolInject.Store(config.DisableAutoToolInject)
149151

150152
manager.logger.Info("%s tool manager initialized with tool execution timeout: %v, max agent depth: %d, and code mode binding level: %s", MCPLogPrefix, config.ToolExecutionTimeout, config.MaxAgentDepth, config.CodeModeBindingLevel)
151153
return manager
@@ -293,6 +295,16 @@ func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schem
293295
return req
294296
}
295297

298+
// When auto tool injection is disabled, only inject tools if the request
299+
// has explicit context filters set (e.g. via x-bf-mcp-include-tools header).
300+
if m.disableAutoToolInject.Load() {
301+
includeTools := ctx.Value(MCPContextKeyIncludeTools)
302+
includeClients := ctx.Value(MCPContextKeyIncludeClients)
303+
if includeTools == nil && includeClients == nil {
304+
return req
305+
}
306+
}
307+
296308
availableTools := m.GetAvailableTools(ctx)
297309

298310
if len(availableTools) == 0 {
@@ -659,9 +671,17 @@ func (m *ToolsManager) UpdateConfig(config *schemas.MCPToolManagerConfig) {
659671
})
660672
}
661673

674+
m.disableAutoToolInject.Store(config.DisableAutoToolInject)
675+
662676
m.logger.Info("%s tool manager configuration updated with tool execution timeout: %v, max agent depth: %d, and code mode binding level: %s", MCPLogPrefix, config.ToolExecutionTimeout, config.MaxAgentDepth, config.CodeModeBindingLevel)
663677
}
664678

679+
// UpdateDisableAutoToolInject updates the auto tool injection toggle atomically.
680+
// This method is safe to call concurrently from multiple goroutines.
681+
func (m *ToolsManager) UpdateDisableAutoToolInject(val bool) {
682+
m.disableAutoToolInject.Store(val)
683+
}
684+
665685
// GetCodeModeBindingLevel returns the current code mode binding level.
666686
// This method is safe to call concurrently from multiple goroutines.
667687
func (m *ToolsManager) GetCodeModeBindingLevel() schemas.CodeModeBindingLevel {

core/schemas/mcp.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,10 @@ type MCPConfig struct {
4747
}
4848

4949
type MCPToolManagerConfig struct {
50-
ToolExecutionTimeout time.Duration `json:"tool_execution_timeout"`
51-
MaxAgentDepth int `json:"max_agent_depth"`
52-
CodeModeBindingLevel CodeModeBindingLevel `json:"code_mode_binding_level,omitempty"` // How tools are exposed in VFS: "server" or "tool"
50+
ToolExecutionTimeout time.Duration `json:"tool_execution_timeout"`
51+
MaxAgentDepth int `json:"max_agent_depth"`
52+
CodeModeBindingLevel CodeModeBindingLevel `json:"code_mode_binding_level,omitempty"` // How tools are exposed in VFS: "server" or "tool"
53+
DisableAutoToolInject bool `json:"disable_auto_tool_inject,omitempty"` // When true, MCP tools are not injected into requests by default
5354
}
5455

5556
const (

framework/changelog.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
- feat: add MCPDisableAutoToolInject column to TableClientConfig

framework/configstore/clientconfig.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,8 @@ type ClientConfig struct {
5454
MCPAgentDepth int `json:"mcp_agent_depth"` // The maximum depth for MCP agent mode tool execution
5555
MCPToolExecutionTimeout int `json:"mcp_tool_execution_timeout"` // The timeout for individual tool execution in seconds
5656
MCPCodeModeBindingLevel string `json:"mcp_code_mode_binding_level"` // Code mode binding level: "server" or "tool"
57-
MCPToolSyncInterval int `json:"mcp_tool_sync_interval"` // Global tool sync interval in minutes (default: 10, 0 = disabled)
57+
MCPToolSyncInterval int `json:"mcp_tool_sync_interval"` // Global tool sync interval in minutes (default: 10, 0 = disabled)
58+
MCPDisableAutoToolInject bool `json:"mcp_disable_auto_tool_inject"` // When true, MCP tools are not injected into requests by default
5859
HeaderFilterConfig *tables.GlobalHeaderFilterConfig `json:"header_filter_config,omitempty"` // Global header filtering configuration for x-bf-eh-* headers
5960
AsyncJobResultTTL int `json:"async_job_result_ttl"` // Default TTL for async job results in seconds (default: 3600 = 1 hour)
6061
RequiredHeaders []string `json:"required_headers,omitempty"` // Headers that must be present on every request (case-insensitive)
@@ -134,6 +135,12 @@ func (c *ClientConfig) GenerateClientConfigHash() (string, error) {
134135
hash.Write([]byte("mcpToolSyncInterval:0"))
135136
}
136137

138+
if c.MCPDisableAutoToolInject {
139+
hash.Write([]byte("mcpDisableAutoToolInject:true"))
140+
} else {
141+
hash.Write([]byte("mcpDisableAutoToolInject:false"))
142+
}
143+
137144
if c.AsyncJobResultTTL > 0 {
138145
hash.Write([]byte("asyncJobResultTTL:" + strconv.Itoa(c.AsyncJobResultTTL)))
139146
} else {

framework/configstore/migrations.go

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
296296
if err := migrationAddBedrockAssumeRoleColumns(ctx, db); err != nil {
297297
return err
298298
}
299+
if err := migrationAddMCPDisableAutoToolInjectColumn(ctx, db); err != nil {
300+
return err
301+
}
299302
return nil
300303
}
301304

@@ -4173,3 +4176,33 @@ func migrationAddBedrockAssumeRoleColumns(ctx context.Context, db *gorm.DB) erro
41734176
}
41744177
return nil
41754178
}
4179+
4180+
// migrationAddMCPDisableAutoToolInjectColumn adds the mcp_disable_auto_tool_inject column to the client config table.
4181+
// When true, MCP tools are not automatically injected into requests; only explicit context filters apply.
4182+
func migrationAddMCPDisableAutoToolInjectColumn(ctx context.Context, db *gorm.DB) error {
4183+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
4184+
ID: "add_mcp_disable_auto_tool_inject_column",
4185+
Migrate: func(tx *gorm.DB) error {
4186+
tx = tx.WithContext(ctx)
4187+
migratorInstance := tx.Migrator()
4188+
if !migratorInstance.HasColumn(&tables.TableClientConfig{}, "mcp_disable_auto_tool_inject") {
4189+
if err := migratorInstance.AddColumn(&tables.TableClientConfig{}, "mcp_disable_auto_tool_inject"); err != nil {
4190+
return err
4191+
}
4192+
}
4193+
return nil
4194+
},
4195+
Rollback: func(tx *gorm.DB) error {
4196+
tx = tx.WithContext(ctx)
4197+
migratorInstance := tx.Migrator()
4198+
if err := migratorInstance.DropColumn(&tables.TableClientConfig{}, "mcp_disable_auto_tool_inject"); err != nil {
4199+
return err
4200+
}
4201+
return nil
4202+
},
4203+
}})
4204+
if err := m.Migrate(); err != nil {
4205+
return fmt.Errorf("error while running mcp disable auto tool inject migration: %s", err.Error())
4206+
}
4207+
return nil
4208+
}

framework/configstore/rdb.go

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,8 @@ func (s *RDBConfigStore) UpdateClientConfig(ctx context.Context, config *ClientC
5656
MCPAgentDepth: config.MCPAgentDepth,
5757
MCPToolExecutionTimeout: config.MCPToolExecutionTimeout,
5858
MCPCodeModeBindingLevel: config.MCPCodeModeBindingLevel,
59-
MCPToolSyncInterval: config.MCPToolSyncInterval,
59+
MCPToolSyncInterval: config.MCPToolSyncInterval,
60+
MCPDisableAutoToolInject: config.MCPDisableAutoToolInject,
6061
AsyncJobResultTTL: config.AsyncJobResultTTL,
6162
RequiredHeaders: config.RequiredHeaders,
6263
LoggingHeaders: config.LoggingHeaders,
@@ -221,8 +222,9 @@ func (s *RDBConfigStore) GetClientConfig(ctx context.Context) (*ClientConfig, er
221222
MCPAgentDepth: dbConfig.MCPAgentDepth,
222223
MCPToolExecutionTimeout: dbConfig.MCPToolExecutionTimeout,
223224
MCPCodeModeBindingLevel: dbConfig.MCPCodeModeBindingLevel,
224-
MCPToolSyncInterval: dbConfig.MCPToolSyncInterval,
225-
AsyncJobResultTTL: dbConfig.AsyncJobResultTTL,
225+
MCPToolSyncInterval: dbConfig.MCPToolSyncInterval,
226+
MCPDisableAutoToolInject: dbConfig.MCPDisableAutoToolInject,
227+
AsyncJobResultTTL: dbConfig.AsyncJobResultTTL,
226228
RequiredHeaders: dbConfig.RequiredHeaders,
227229
LoggingHeaders: dbConfig.LoggingHeaders,
228230
HeaderFilterConfig: dbConfig.HeaderFilterConfig,
@@ -894,9 +896,10 @@ func (s *RDBConfigStore) GetMCPConfig(ctx context.Context) (*schemas.MCPConfig,
894896
return nil, err
895897
}
896898
toolManagerConfig := schemas.MCPToolManagerConfig{
897-
ToolExecutionTimeout: time.Duration(clientConfig.MCPToolExecutionTimeout) * time.Second,
898-
MaxAgentDepth: clientConfig.MCPAgentDepth,
899-
CodeModeBindingLevel: schemas.CodeModeBindingLevel(clientConfig.MCPCodeModeBindingLevel),
899+
ToolExecutionTimeout: time.Duration(clientConfig.MCPToolExecutionTimeout) * time.Second,
900+
MaxAgentDepth: clientConfig.MCPAgentDepth,
901+
CodeModeBindingLevel: schemas.CodeModeBindingLevel(clientConfig.MCPCodeModeBindingLevel),
902+
DisableAutoToolInject: clientConfig.MCPDisableAutoToolInject,
900903
}
901904
clientConfigs := make([]*schemas.MCPClientConfig, len(dbMCPClients))
902905
for i, dbClient := range dbMCPClients {

0 commit comments

Comments
 (0)