Skip to content

Commit 8c5bb95

Browse files
Pratham-Mishra04akshaydeo
authored andcommitted
feat: plugin schema extensions for mcp plugins
1 parent 0288181 commit 8c5bb95

File tree

157 files changed

+14143
-3666
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

157 files changed

+14143
-3666
lines changed

core/bifrost.go

Lines changed: 637 additions & 159 deletions
Large diffs are not rendered by default.

core/chatbot_test.go

Lines changed: 0 additions & 941 deletions
This file was deleted.

core/internal/testutil/setup.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@ func getBifrost(ctx context.Context) (*bifrost.Bifrost, error) {
3737
// Initialize Bifrost
3838
b, err := bifrost.Init(ctx, schemas.BifrostConfig{
3939
Account: &account,
40-
Plugins: nil,
4140
Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug),
4241
})
4342
if err != nil {

core/mcp/agent.go

Lines changed: 24 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import (
2020
// - initialResponse: The initial chat response containing tool calls
2121
// - makeReq: Function to make subsequent chat requests during agent execution
2222
// - fetchNewRequestIDFunc: Optional function to generate unique request IDs for each iteration
23-
// - executeToolFunc: Function to execute individual tool calls
23+
// - executeToolFunc: Function to execute individual tool calls using unified MCP request/response
2424
// - clientManager: Client manager for accessing MCP clients and tools
2525
//
2626
// Returns:
@@ -33,7 +33,7 @@ func ExecuteAgentForChatRequest(
3333
initialResponse *schemas.BifrostChatResponse,
3434
makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError),
3535
fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string,
36-
executeToolFunc func(ctx *schemas.BifrostContext, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, error),
36+
executeToolFunc func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error),
3737
clientManager ClientManager,
3838
) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
3939
// Create adapter for Chat API
@@ -73,7 +73,7 @@ func ExecuteAgentForChatRequest(
7373
// - initialResponse: The initial responses response containing tool calls
7474
// - makeReq: Function to make subsequent responses requests during agent execution
7575
// - fetchNewRequestIDFunc: Optional function to generate unique request IDs for each iteration
76-
// - executeToolFunc: Function to execute individual tool calls
76+
// - executeToolFunc: Function to execute individual tool calls using unified MCP request/response
7777
// - clientManager: Client manager for accessing MCP clients and tools
7878
//
7979
// Returns:
@@ -86,7 +86,7 @@ func ExecuteAgentForResponsesRequest(
8686
initialResponse *schemas.BifrostResponsesResponse,
8787
makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError),
8888
fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string,
89-
executeToolFunc func(ctx *schemas.BifrostContext, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, error),
89+
executeToolFunc func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error),
9090
clientManager ClientManager,
9191
) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
9292
// Create adapter for Responses API
@@ -125,7 +125,7 @@ func ExecuteAgentForResponsesRequest(
125125
// - maxAgentDepth: Maximum number of agent iterations allowed
126126
// - adapter: API adapter that abstracts differences between Chat and Responses APIs
127127
// - fetchNewRequestIDFunc: Optional function to generate unique request IDs for each iteration
128-
// - executeToolFunc: Function to execute individual tool calls
128+
// - executeToolFunc: Function to execute individual tool calls using unified MCP request/response
129129
// - clientManager: Client manager for accessing MCP clients and tools
130130
//
131131
// Returns:
@@ -136,7 +136,7 @@ func executeAgent(
136136
maxAgentDepth int,
137137
adapter agentAPIAdapter,
138138
fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string,
139-
executeToolFunc func(ctx *schemas.BifrostContext, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, error),
139+
executeToolFunc func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error),
140140
clientManager ClientManager,
141141
) (interface{}, *schemas.BifrostError) {
142142
logger.Debug("Entering agent mode - detected tool calls in response")
@@ -292,12 +292,24 @@ func executeAgent(
292292
for _, toolCall := range autoExecutableTools {
293293
go func(toolCall schemas.ChatAssistantMessageToolCall) {
294294
defer wg.Done()
295-
toolResult, toolErr := executeToolFunc(ctx, toolCall)
295+
// Create MCP request for this tool call
296+
mcpRequest := &schemas.BifrostMCPRequest{
297+
RequestType: schemas.MCPRequestTypeChatToolCall,
298+
ChatAssistantMessageToolCall: &toolCall,
299+
}
300+
301+
mcpResponse, toolErr := executeToolFunc(ctx, mcpRequest)
296302
if toolErr != nil {
297-
logger.Warn("Tool execution failed: %v", toolErr)
303+
logger.Warn(fmt.Sprintf("Tool execution failed: %v", toolErr))
298304
channelToolResults <- createToolResultMessage(toolCall, "", toolErr)
305+
} else if mcpResponse != nil && mcpResponse.ChatMessage != nil {
306+
channelToolResults <- mcpResponse.ChatMessage
307+
} else if mcpResponse != nil && mcpResponse.ChatMessage == nil {
308+
// Send empty result when mcpResponse is non-nil but ChatMessage is nil
309+
channelToolResults <- createToolResultMessage(toolCall, "", nil)
299310
} else {
300-
channelToolResults <- toolResult
311+
// Fallback: send empty result when both mcpResponse and toolErr are nil
312+
channelToolResults <- createToolResultMessage(toolCall, "", nil)
301313
}
302314
}(toolCall)
303315
}
@@ -457,8 +469,9 @@ func buildAllowedAutoExecutionTools(ctx *schemas.BifrostContext, clientManager C
457469
autoExecutableTools = append(autoExecutableTools, "*")
458470
continue
459471
}
460-
// Use parsed tool name (as it appears in code)
461-
parsedToolName := parseToolName(originalToolName)
472+
// Replace - with _ for code mode compatibility, then parse for JS compatibility
473+
toolNameForCode := strings.ReplaceAll(originalToolName, "-", "_")
474+
parsedToolName := parseToolName(toolNameForCode)
462475
autoExecutableTools = append(autoExecutableTools, parsedToolName)
463476
}
464477

core/mcp/clientmanager.go

Lines changed: 68 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -189,8 +189,7 @@ func (m *MCPManager) EditClient(id string, updatedConfig schemas.MCPClientConfig
189189
return fmt.Errorf("invalid MCP client configuration: %w", err)
190190
}
191191

192-
// Check if is_ping_available changed
193-
isPingAvailableChanged := client.ExecutionConfig.IsPingAvailable != updatedConfig.IsPingAvailable
192+
oldName := client.ExecutionConfig.Name
194193

195194
// Update the client's execution config with new tool filters
196195
config := client.ExecutionConfig
@@ -199,17 +198,42 @@ func (m *MCPManager) EditClient(id string, updatedConfig schemas.MCPClientConfig
199198
config.Headers = updatedConfig.Headers
200199
config.ToolsToExecute = updatedConfig.ToolsToExecute
201200
config.ToolsToAutoExecute = updatedConfig.ToolsToAutoExecute
202-
config.IsPingAvailable = updatedConfig.IsPingAvailable
201+
config.IsCodeModeClient = updatedConfig.IsCodeModeClient
203202

204203
// Store the updated config
205204
client.ExecutionConfig = config
206205

207-
// If is_ping_available changed, update the health monitor
208-
if isPingAvailableChanged {
209-
// Stop and restart the health monitor with the new is_ping_available setting
210-
m.healthMonitorManager.StopMonitoring(id)
211-
monitor := NewClientHealthMonitor(m, id, DefaultHealthCheckInterval, config.IsPingAvailable)
212-
m.healthMonitorManager.StartMonitoring(monitor)
206+
// If the client name has changed, update all tool name prefixes in the ToolMap
207+
if oldName != updatedConfig.Name {
208+
oldPrefix := oldName + "-"
209+
newPrefix := updatedConfig.Name + "-"
210+
211+
// Create a new ToolMap with updated tool names
212+
newToolMap := make(map[string]schemas.ChatTool, len(client.ToolMap))
213+
for oldToolName, tool := range client.ToolMap {
214+
var newToolName string
215+
if strings.HasPrefix(oldToolName, oldPrefix) {
216+
// Update the tool name by replacing the old prefix with the new prefix
217+
newToolName = newPrefix + strings.TrimPrefix(oldToolName, oldPrefix)
218+
} else {
219+
newToolName = oldToolName
220+
}
221+
222+
// Update the tool's function name if it's a function tool
223+
if tool.Function != nil {
224+
updatedTool := tool
225+
updatedTool.Function.Name = newToolName
226+
newToolMap[newToolName] = updatedTool
227+
} else {
228+
newToolMap[newToolName] = tool
229+
}
230+
}
231+
232+
// Replace the old ToolMap with the new one
233+
client.ToolMap = newToolMap
234+
235+
// Also update the client Name field
236+
client.Name = updatedConfig.Name
213237
}
214238

215239
return nil
@@ -339,6 +363,7 @@ func (m *MCPManager) connectToMCPClient(config schemas.MCPClientConfig) error {
339363
var err error
340364

341365
// Create appropriate transport based on connection type
366+
logger.Debug(fmt.Sprintf("%s [%s] Creating %s connection...", MCPLogPrefix, config.Name, config.ConnectionType))
342367
switch config.ConnectionType {
343368
case schemas.MCPConnectionTypeHTTP:
344369
externalClient, connectionInfo, err = m.createHTTPConnection(config)
@@ -355,29 +380,39 @@ func (m *MCPManager) connectToMCPClient(config schemas.MCPClientConfig) error {
355380
if err != nil {
356381
return fmt.Errorf("failed to create connection: %w", err)
357382
}
383+
logger.Debug(fmt.Sprintf("%s [%s] Connection created successfully", MCPLogPrefix, config.Name))
358384

359385
// Initialize the external client with timeout
360-
// For SSE connections, we need a long-lived context, for others we can use timeout
386+
// For SSE and STDIO connections, we need a long-lived context for the connection
387+
// but use a timeout context for the initialization phase to prevent indefinite hangs
361388
var ctx context.Context
362389
var cancel context.CancelFunc
390+
var longLivedCtx context.Context
391+
var longLivedCancel context.CancelFunc
392+
393+
if config.ConnectionType == schemas.MCPConnectionTypeSSE || config.ConnectionType == schemas.MCPConnectionTypeSTDIO {
394+
// Create long-lived context for the connection (subprocess lifetime)
395+
longLivedCtx, longLivedCancel = context.WithCancel(m.ctx)
363396

364-
if config.ConnectionType == schemas.MCPConnectionTypeSSE {
365-
// SSE connections need a long-lived context for the persistent stream
366-
ctx, cancel = context.WithCancel(m.ctx)
367-
// Don't defer cancel here - SSE needs the context to remain active
397+
// Use long-lived context for starting the transport (spawns subprocess)
398+
// but create a timeout context for initialization to prevent hangs
399+
ctx = longLivedCtx
400+
cancel = longLivedCancel
368401
} else {
369402
// Other connection types can use timeout context
370403
ctx, cancel = context.WithTimeout(m.ctx, MCPClientConnectionEstablishTimeout)
371404
defer cancel()
372405
}
373406

374407
// Start the transport first (required for STDIO and SSE clients)
408+
logger.Debug(fmt.Sprintf("%s [%s] Starting transport...", MCPLogPrefix, config.Name))
375409
if err := externalClient.Start(ctx); err != nil {
376410
if config.ConnectionType == schemas.MCPConnectionTypeSSE {
377411
cancel() // Cancel SSE context only on error
378412
}
379413
return fmt.Errorf("failed to start MCP client transport %s: %v", config.Name, err)
380414
}
415+
logger.Debug(fmt.Sprintf("%s [%s] Transport started successfully", MCPLogPrefix, config.Name))
381416

382417
// Create proper initialize request for external client
383418
extInitRequest := mcp.InitializeRequest{
@@ -391,21 +426,39 @@ func (m *MCPManager) connectToMCPClient(config schemas.MCPClientConfig) error {
391426
},
392427
}
393428

394-
_, err = externalClient.Initialize(ctx, extInitRequest)
429+
// For STDIO/SSE: Use a timeout context for initialization to prevent indefinite hangs
430+
// The subprocess will continue running with the long-lived context
431+
var initCtx context.Context
432+
var initCancel context.CancelFunc
433+
434+
if config.ConnectionType == schemas.MCPConnectionTypeSSE || config.ConnectionType == schemas.MCPConnectionTypeSTDIO {
435+
// Create timeout context for initialization phase only
436+
initCtx, initCancel = context.WithTimeout(longLivedCtx, MCPClientConnectionEstablishTimeout)
437+
defer initCancel()
438+
logger.Debug(fmt.Sprintf("%s [%s] Initializing client with %v timeout...", MCPLogPrefix, config.Name, MCPClientConnectionEstablishTimeout))
439+
} else {
440+
// HTTP already has timeout
441+
initCtx = ctx
442+
}
443+
444+
_, err = externalClient.Initialize(initCtx, extInitRequest)
395445
if err != nil {
396446
if config.ConnectionType == schemas.MCPConnectionTypeSSE {
397447
cancel() // Cancel SSE context only on error
398448
}
399449
return fmt.Errorf("failed to initialize MCP client %s: %v", config.Name, err)
400450
}
451+
logger.Debug(fmt.Sprintf("%s [%s] Client initialized successfully", MCPLogPrefix, config.Name))
401452

402453
// Retrieve tools from the external server (this also requires network I/O)
454+
logger.Debug(fmt.Sprintf("%s [%s] Retrieving tools...", MCPLogPrefix, config.Name))
403455
tools, err := retrieveExternalTools(ctx, externalClient, config.Name)
404456
if err != nil {
405457
logger.Warn("%s Failed to retrieve tools from %s: %v", MCPLogPrefix, config.Name, err)
406458
// Continue with connection even if tool retrieval fails
407459
tools = make(map[string]schemas.ChatTool)
408460
}
461+
logger.Debug(fmt.Sprintf("%s [%s] Retrieved %d tools", MCPLogPrefix, config.Name, len(tools)))
409462

410463
// Second lock: Update client with final connection details and tools
411464
m.mu.Lock()

0 commit comments

Comments
 (0)