From 576df36266078b8c603761e001a54c38da8db023 Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Mon, 11 May 2026 21:36:49 +0530 Subject: [PATCH] feat: extend mcp plugin interface for list tools, ping and connections --- core/bifrost.go | 437 +++++-------- core/internal/mcptests/agent_adapter_test.go | 27 - core/internal/mcptests/agent_basic_test.go | 28 - .../mcptests/agent_context_filtering_test.go | 24 - .../mcptests/agent_error_handling_test.go | 18 - .../internal/mcptests/agent_filtering_test.go | 39 -- core/internal/mcptests/agent_limits_test.go | 58 +- .../mcptests/agent_mixed_permissions_test.go | 24 - .../mcptests/agent_multiconnection_test.go | 15 - .../mcptests/agent_parallel_execution_test.go | 18 - .../mcptests/agent_request_id_test.go | 18 - .../mcptests/agent_state_transitions_test.go | 21 - core/internal/mcptests/agent_test_helpers.go | 6 - .../agent_test_helpers_example_test.go | 12 - .../mcptests/codemode_agent_multiturn_test.go | 18 - .../codemode_agent_singleturn_test.go | 6 - core/internal/mcptests/codemode_agent_test.go | 45 -- .../mcptests/codemode_vs_noncodemode_test.go | 9 - .../mcptests/connect_ping_listtools_test.go | 599 ++++++++++++++++++ .../mcptests/context_propagation_test.go | 12 - core/internal/mcptests/test_plugins.go | 477 +++++++++++++- core/internal/mcptests/tool_call_id_test.go | 6 - core/mcp/agent.go | 6 +- core/mcp/agent_test.go | 6 + core/mcp/clientmanager.go | 531 ++++++++++------ core/mcp/codemode.go | 6 - core/mcp/codemode/starlark/executecode.go | 11 +- core/mcp/codemode/starlark/starlark.go | 10 +- core/mcp/codemode/starlark/starlark_test.go | 6 + core/mcp/exec.go | 200 ++++++ core/mcp/healthmonitor.go | 20 +- core/mcp/interface.go | 15 +- core/mcp/mcp.go | 39 +- core/mcp/pluginpipeline.go | 378 +++++++++++ core/mcp/toolmanager.go | 110 +--- core/mcp/toolmanager_test.go | 15 +- core/mcp/toolsync.go | 2 +- core/mcp/utils.go | 38 +- core/schemas/bifrost.go | 232 ++++++- core/schemas/plugin.go | 41 ++ core/schemas/plugin_native.go | 8 + core/schemas/trace.go | 3 + core/utils.go | 38 +- framework/plugins/soloader.go | 17 + framework/plugins/soplugin.go | 27 + plugins/governance/main.go | 18 + plugins/logging/main.go | 19 +- ui/app/workspace/config/views/mcpView.tsx | 2 +- 48 files changed, 2646 insertions(+), 1069 deletions(-) create mode 100644 core/internal/mcptests/connect_ping_listtools_test.go create mode 100644 core/mcp/exec.go create mode 100644 core/mcp/pluginpipeline.go diff --git a/core/bifrost.go b/core/bifrost.go index 0c5eaf5445a..5fc49bef91b 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -5,7 +5,6 @@ package bifrost import ( "context" - "errors" "fmt" "slices" "sort" @@ -79,7 +78,6 @@ type Bifrost struct { responseStreamPool sync.Pool // Pool for response stream channels, initial pool size is set in Init pluginPipelinePool sync.Pool // Pool for PluginPipeline objects bifrostRequestPool sync.Pool // Pool for BifrostRequest objects - mcpRequestPool sync.Pool // Pool for BifrostMCPRequest objects 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) @@ -291,12 +289,8 @@ func Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error) { return &schemas.BifrostRequest{} }, } - bifrost.mcpRequestPool = sync.Pool{ - New: func() interface{} { - return &schemas.BifrostMCPRequest{} - }, - } - // Prewarm pools with multiple objects + // Prewarm pools. The MCP request pool is owned by the mcp package now — + // see core/mcp/exec.go. for range config.InitialPoolSize { // Create and put new objects directly into pools bifrost.channelMessagePool.Put(&ChannelMessage{}) @@ -308,7 +302,6 @@ func Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error) { postHookErrors: make([]error, 0), }) bifrost.bifrostRequestPool.Put(&schemas.BifrostRequest{}) - bifrost.mcpRequestPool.Put(&schemas.BifrostMCPRequest{}) } providerKeys, err := bifrost.account.GetConfiguredProviders() @@ -523,7 +516,7 @@ func (bifrost *Bifrost) ListAllModels(ctx *schemas.BifrostContext, req *schemas. if !strings.Contains(bifrostErr.Error.Message, "no keys found") && !strings.Contains(bifrostErr.Error.Message, "not supported") { providerErr = bifrostErr - bifrost.logger.Warn("failed to list models for provider %s: %s", providerKey, GetErrorMessage(bifrostErr)) + bifrost.logger.Warn("failed to list models for provider %s: %s", providerKey, bifrostErr.GetErrorString()) } // Collect key statuses from error (failure case) if len(bifrostErr.ExtraFields.KeyStatuses) > 0 { @@ -732,14 +725,13 @@ func (bifrost *Bifrost) ChatCompletionRequest(ctx *schemas.BifrostContext, req * return nil, err } - // Check if we should enter agent mode + // Check if we should enter agent mode. if bifrost.MCPManager != nil { return bifrost.MCPManager.CheckAndExecuteAgentForChatRequest( ctx, req, response, bifrost.makeChatCompletionRequest, - bifrost.executeMCPToolWithHooks, ) } @@ -835,14 +827,13 @@ func (bifrost *Bifrost) ResponsesRequest(ctx *schemas.BifrostContext, req *schem return nil, err } - // Check if we should enter agent mode + // Check if we should enter agent mode. if bifrost.MCPManager != nil { return bifrost.MCPManager.CheckAndExecuteAgentForResponsesRequest( ctx, req, response, bifrost.makeResponsesRequest, - bifrost.executeMCPToolWithHooks, ) } @@ -2496,117 +2487,37 @@ func (bifrost *Bifrost) PassthroughStream( } // ExecuteChatMCPTool executes an MCP tool call and returns the result as a chat message. -// This is the main public API for manual MCP tool execution in Chat format. -// -// Parameters: -// - ctx: Execution context -// - toolCall: The tool call to execute (from assistant message) -// -// Returns: -// - *schemas.ChatMessage: Tool message with execution result -// - *schemas.BifrostError: Any execution error +// This is the main public API for manual MCP tool execution in Chat format. All the +// real work — request pooling, plugin gate (PreMCPHook / PostMCPHook), short-circuit +// handling, error enrichment — lives on MCPManager.ExecuteChatTool. func (bifrost *Bifrost) ExecuteChatMCPTool(ctx *schemas.BifrostContext, toolCall *schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, *schemas.BifrostError) { - // Handle nil context early to prevent issues downstream if ctx == nil { ctx = bifrost.ctx } - - // Validate toolCall is not nil - if toolCall == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "toolCall cannot be nil", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: schemas.ChatCompletionRequest, - }, - } - } - - // Get MCP request from pool and populate - mcpRequest := bifrost.getMCPRequest() - mcpRequest.RequestType = schemas.MCPRequestTypeChatToolCall - mcpRequest.ChatAssistantMessageToolCall = toolCall - defer bifrost.releaseMCPRequest(mcpRequest) - - // Execute with common handler - result, err := bifrost.handleMCPToolExecution(ctx, mcpRequest, schemas.ChatCompletionRequest) - if err != nil { - return nil, err - } - - // Validate and extract chat message from result - if result == nil || result.ChatMessage == nil { + if bifrost.MCPManager == nil { return nil, &schemas.BifrostError{ IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "MCP tool execution returned nil chat message", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: schemas.ChatCompletionRequest, - }, + Error: &schemas.ErrorField{Message: "mcp is not configured in this bifrost instance"}, + ExtraFields: schemas.BifrostErrorExtraFields{RequestType: schemas.ChatCompletionRequest}, } } - - return result.ChatMessage, nil + return bifrost.MCPManager.ExecuteChatTool(ctx, toolCall) } -// ExecuteResponsesMCPTool executes an MCP tool call and returns the result as a responses message. -// This is the main public API for manual MCP tool execution in Responses format. -// -// Parameters: -// - ctx: Execution context -// - toolCall: The tool call to execute (from assistant message) -// -// Returns: -// - *schemas.ResponsesMessage: Tool message with execution result -// - *schemas.BifrostError: Any execution error +// ExecuteResponsesMCPTool executes an MCP tool call and returns the result as a responses +// message. Thin delegator — see ExecuteChatMCPTool for the rationale. func (bifrost *Bifrost) ExecuteResponsesMCPTool(ctx *schemas.BifrostContext, toolCall *schemas.ResponsesToolMessage) (*schemas.ResponsesMessage, *schemas.BifrostError) { - // Handle nil context early to prevent issues downstream if ctx == nil { ctx = bifrost.ctx } - - // Validate toolCall is not nil - if toolCall == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "toolCall cannot be nil", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: schemas.ResponsesRequest, - }, - } - } - - // Get MCP request from pool and populate - mcpRequest := bifrost.getMCPRequest() - mcpRequest.RequestType = schemas.MCPRequestTypeResponsesToolCall - mcpRequest.ResponsesToolMessage = toolCall - defer bifrost.releaseMCPRequest(mcpRequest) - - // Execute with common handler - result, err := bifrost.handleMCPToolExecution(ctx, mcpRequest, schemas.ResponsesRequest) - if err != nil { - return nil, err - } - - // Validate and extract responses message from result - if result == nil || result.ResponsesMessage == nil { + if bifrost.MCPManager == nil { return nil, &schemas.BifrostError{ IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "MCP tool execution returned nil responses message", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: schemas.ResponsesRequest, - }, + Error: &schemas.ErrorField{Message: "mcp is not configured in this bifrost instance"}, + ExtraFields: schemas.BifrostErrorExtraFields{RequestType: schemas.ResponsesRequest}, } } - - return result.ResponsesMessage, nil + return bifrost.MCPManager.ExecuteResponsesTool(ctx, toolCall) } // ContainerCreateRequest creates a new container. @@ -5548,7 +5459,7 @@ func executeRequestWithRetries[T any]( (bifrostError.Error.Type != nil && IsRateLimitErrorMessage(*bifrostError.Error.Type)) || (bifrostError.Error.Code != nil && IsRateLimitErrorMessage(*bifrostError.Error.Code)))) - errMessage := GetErrorMessage(bifrostError) + errMessage := bifrostError.GetErrorString() if bifrostError.Error != nil && (bifrostError.Error.Message == schemas.ErrProviderDoRequest || @@ -6369,166 +6280,6 @@ func (bifrost *Bifrost) handleProviderStreamRequest(provider schemas.Provider, r } } -// handleMCPToolExecution is the common handler for MCP tool execution with plugin pipeline support. -// It handles pre-hooks, execution, post-hooks, and error handling for both Chat and Responses formats. -// -// Parameters: -// - ctx: Execution context -// - mcpRequest: The MCP request to execute (already populated with tool call) -// - requestType: The request type for error reporting (ChatCompletionRequest or ResponsesRequest) -// -// Returns: -// - *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 { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "mcp is not configured in this bifrost instance", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: requestType, - }, - } - } - - // Ensure request ID exists for hooks/tracing consistency - if _, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string); !ok { - ctx.SetValue(schemas.BifrostContextKeyRequestID, uuid.New().String()) - } - - // Get plugin pipeline for MCP hooks - pipeline := bifrost.getPluginPipeline() - defer bifrost.releasePluginPipeline(pipeline) - - // Run pre-hooks - preReq, shortCircuit, preCount := pipeline.RunMCPPreHooks(ctx, mcpRequest) - - // Handle short-circuit cases - if shortCircuit != nil { - // Handle short-circuit with response (success case) - if shortCircuit.Response != nil { - finalMcpResp, bifrostErr := pipeline.RunMCPPostHooks(ctx, shortCircuit.Response, nil, preCount) - drainAndAttachPluginLogs(ctx) - if bifrostErr != nil { - return nil, bifrostErr - } - return finalMcpResp, nil - } - // Handle short-circuit with error - if shortCircuit.Error != nil { - // Capture post-hook results to respect transformations or recovery - finalResp, finalErr := pipeline.RunMCPPostHooks(ctx, nil, shortCircuit.Error, preCount) - drainAndAttachPluginLogs(ctx) - // Return post-hook error if present (post-hook may have transformed the error) - if finalErr != nil { - return nil, finalErr - } - // Return post-hook response if present (post-hook may have recovered from error) - if finalResp != nil { - return finalResp, nil - } - // Fall back to original short-circuit error if post-hooks returned nil/nil - return nil, shortCircuit.Error - } - } - - if preReq == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "MCP request after plugin hooks cannot be nil", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: requestType, - }, - } - } - - // Execute tool with modified request - result, err := bifrost.MCPManager.ExecuteToolCall(ctx, preReq) - - // Prepare MCP response and error for post-hooks - var mcpResp *schemas.BifrostMCPResponse - var bifrostErr *schemas.BifrostError - - if err != nil { - bifrostErr = &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: err.Error(), - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: requestType, - }, - } - // Preserve MCPUserOAuthRequiredError for downstream detection in agent mode - var oauthErr *schemas.MCPUserOAuthRequiredError - if errors.As(err, &oauthErr) { - bifrostErr.ExtraFields.MCPAuthRequired = oauthErr - } - } else if result == nil { - bifrostErr = &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "tool execution returned nil result", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: requestType, - }, - } - } else { - // Use the MCP response directly - mcpResp = result - } - - // Run post-hooks - finalResp, finalErr := pipeline.RunMCPPostHooks(ctx, mcpResp, bifrostErr, preCount) - drainAndAttachPluginLogs(ctx) - - if finalErr != nil { - return nil, finalErr - } - - return finalResp, nil -} - -// executeMCPToolWithHooks is a wrapper around handleMCPToolExecution that matches the signature -// expected by the agent's executeToolFunc parameter. It runs MCP plugin hooks before and after -// tool execution to enable logging, telemetry, and other plugin functionality. -func (bifrost *Bifrost) executeMCPToolWithHooks(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - // Defensive check: context must be non-nil to prevent panics in plugin hooks - if ctx == nil { - return nil, fmt.Errorf("context cannot be nil") - } - - if request == nil { - return nil, fmt.Errorf("request cannot be nil") - } - - // Determine request type from the MCP request - explicitly handle all known types - var requestType schemas.RequestType - switch request.RequestType { - case schemas.MCPRequestTypeChatToolCall: - requestType = schemas.ChatCompletionRequest - case schemas.MCPRequestTypeResponsesToolCall: - requestType = schemas.ResponsesRequest - default: - // Return error for unknown/unsupported request types instead of silently defaulting - return nil, fmt.Errorf("unsupported MCP request type: %s", request.RequestType) - } - - resp, bifrostErr := bifrost.handleMCPToolExecution(ctx, request, requestType) - if bifrostErr != nil { - if bifrostErr.ExtraFields.MCPAuthRequired != nil { - return nil, bifrostErr.ExtraFields.MCPAuthRequired - } - return nil, fmt.Errorf("%s", GetErrorMessage(bifrostErr)) - } - return resp, nil -} - // PLUGIN MANAGEMENT // RunLLMPreHooks executes PreHooks in order, tracks how many ran, and returns the final request, any short-circuit decision, and the count. @@ -6670,8 +6421,9 @@ func (p *PluginPipeline) RunPostLLMHooks(ctx *schemas.BifrostContext, resp *sche } // RunMCPPreHooks executes MCP PreHooks in order for all registered MCP plugins. -// Returns the modified request, any short-circuit decision, and the count of hooks that ran. -// If a plugin short-circuits, only PostHooks for plugins up to and including that plugin will run. +// Handles the envelope-based MCP pipeline (Ping / ListTools / ExecuteTool variants). +// Connect requests do NOT flow through here — they use RunMCPPreConnectionHooks +// with typed signatures. func (p *PluginPipeline) RunMCPPreHooks(ctx *schemas.BifrostContext, req *schemas.BifrostMCPRequest) (*schemas.BifrostMCPRequest, *schemas.MCPPluginShortCircuit, int) { // If the skip plugin pipeline flag is set, skip the plugin pipeline if skipPluginPipeline, ok := ctx.Value(schemas.BifrostContextKeySkipPluginPipeline).(bool); ok && skipPluginPipeline { @@ -6718,10 +6470,9 @@ func (p *PluginPipeline) RunMCPPreHooks(ctx *schemas.BifrostContext, req *schema return req, nil, p.executedPreHooks } -// RunMCPPostHooks executes MCP PostHooks in reverse order for the plugins whose PreMCPHook ran. -// Accepts the MCP response and error, and allows plugins to transform either (e.g., recover from error, or invalidate a response). -// Returns the final MCP response and error after all hooks. If both are set, error takes precedence unless error is nil. -// runFrom is the count of plugins whose PreHooks ran; PostHooks will run in reverse from index (runFrom - 1) down to 0 +// RunMCPPostHooks executes MCP PostHooks in reverse order for the envelope-based +// pipeline (Ping / ListTools / ExecuteTool variants). Connect responses do NOT +// flow through here — they use RunMCPPostConnectionHooks. func (p *PluginPipeline) RunMCPPostHooks(ctx *schemas.BifrostContext, mcpResp *schemas.BifrostMCPResponse, bifrostErr *schemas.BifrostError, runFrom int) (*schemas.BifrostMCPResponse, *schemas.BifrostError) { // If the skip plugin pipeline flag is set, skip the plugin pipeline if skipPluginPipeline, ok := ctx.Value(schemas.BifrostContextKeySkipPluginPipeline).(bool); ok && skipPluginPipeline { @@ -6778,6 +6529,125 @@ func (p *PluginPipeline) RunMCPPostHooks(ctx *schemas.BifrostContext, mcpResp *s return mcpResp, nil } +// RunMCPPreConnectionHooks executes typed Connect PreHooks in order for plugins +// implementing MCPConnectionPlugin. Plugins that only implement MCPPlugin (no typed +// Connect methods) are silently skipped — they cannot observe or intercept the +// connection lifecycle. +// +// Returns the (possibly mutated) typed sub-request, any short-circuit decision, and +// the count of hooks that executed (for matching PostHook dispatch). +func (p *PluginPipeline) RunMCPPreConnectionHooks(ctx *schemas.BifrostContext, req *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectRequest, *schemas.MCPConnectionShortCircuit, int) { + if skipPluginPipeline, ok := ctx.Value(schemas.BifrostContextKeySkipPluginPipeline).(bool); ok && skipPluginPipeline { + return req, nil, 0 + } + var shortCircuit *schemas.MCPConnectionShortCircuit + var err error + ctx.BlockRestrictedWrites() + defer ctx.UnblockRestrictedWrites() + for i, plugin := range p.mcpPlugins { + pluginName := plugin.GetName() + p.logger.Debug("running MCP connect pre-hook for plugin %s", pluginName) + spanCtx, handle := p.tracer.StartSpan(ctx, fmt.Sprintf("plugin.%s.mcp_connect_prehook", sanitizeSpanName(pluginName)), schemas.SpanKindPlugin) + if spanCtx != nil { + if spanID, ok := spanCtx.Value(schemas.BifrostContextKeySpanID).(string); ok { + ctx.SetValue(schemas.BifrostContextKeySpanID, spanID) + } + } + + pluginCtx := ctx.WithPluginScope(&pluginName) + shortCircuit = nil + err = nil + + if cp, ok := plugin.(schemas.MCPConnectionPlugin); ok { + req, shortCircuit, err = cp.PreMCPConnectionHook(pluginCtx, req) + } else { + // Plugin only implements MCPPlugin — Connect is invisible to it. + pluginCtx.ReleasePluginScope() + p.tracer.EndSpan(handle, schemas.SpanStatusOk, "skipped (not MCPConnectionPlugin)") + p.executedPreHooks = i + 1 + continue + } + + pluginCtx.ReleasePluginScope() + + if err != nil { + p.tracer.SetAttribute(handle, "error", err.Error()) + p.tracer.EndSpan(handle, schemas.SpanStatusError, err.Error()) + p.preHookErrors = append(p.preHookErrors, err) + p.logger.Warn("error in PreMCPConnectionHook for plugin %s: %s", pluginName, err.Error()) + } else if shortCircuit != nil { + p.tracer.SetAttribute(handle, "short_circuit", true) + p.tracer.EndSpan(handle, schemas.SpanStatusOk, "short-circuit") + } else { + p.tracer.EndSpan(handle, schemas.SpanStatusOk, "") + } + + p.executedPreHooks = i + 1 + if shortCircuit != nil { + return req, shortCircuit, p.executedPreHooks + } + } + return req, nil, p.executedPreHooks +} + +// RunMCPPostConnectionHooks executes typed Connect PostHooks in reverse order for +// the plugins whose PreMCPConnectionHook ran. Plugins that only implement MCPPlugin +// are skipped (they didn't run in PreHook, they don't run in PostHook). +func (p *PluginPipeline) RunMCPPostConnectionHooks(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPConnectResponse, bifrostErr *schemas.BifrostError, runFrom int) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError) { + if skipPluginPipeline, ok := ctx.Value(schemas.BifrostContextKeySkipPluginPipeline).(bool); ok && skipPluginPipeline { + return resp, bifrostErr + } + if runFrom < 0 { + runFrom = 0 + } + if runFrom > len(p.mcpPlugins) { + runFrom = len(p.mcpPlugins) + } + ctx.BlockRestrictedWrites() + defer ctx.UnblockRestrictedWrites() + var err error + for i := runFrom - 1; i >= 0; i-- { + plugin := p.mcpPlugins[i] + pluginName := plugin.GetName() + p.logger.Debug("running MCP connect post-hook for plugin %s", pluginName) + spanCtx, handle := p.tracer.StartSpan(ctx, fmt.Sprintf("plugin.%s.mcp_connect_posthook", sanitizeSpanName(pluginName)), schemas.SpanKindPlugin) + if spanCtx != nil { + if spanID, ok := spanCtx.Value(schemas.BifrostContextKeySpanID).(string); ok { + ctx.SetValue(schemas.BifrostContextKeySpanID, spanID) + } + } + + pluginCtx := ctx.WithPluginScope(&pluginName) + err = nil + + cp, ok := plugin.(schemas.MCPConnectionPlugin) + if !ok { + pluginCtx.ReleasePluginScope() + p.tracer.EndSpan(handle, schemas.SpanStatusOk, "skipped (not MCPConnectionPlugin)") + continue + } + resp, bifrostErr, err = cp.PostMCPConnectionHook(pluginCtx, resp, bifrostErr) + pluginCtx.ReleasePluginScope() + + if err != nil { + p.tracer.SetAttribute(handle, "error", err.Error()) + p.tracer.EndSpan(handle, schemas.SpanStatusError, err.Error()) + p.postHookErrors = append(p.postHookErrors, err) + p.logger.Warn("error in PostMCPConnectionHook for plugin %s: %v", pluginName, err) + } else { + p.tracer.EndSpan(handle, schemas.SpanStatusOk, "") + } + } + if bifrostErr != nil { + if resp != nil && bifrostErr.StatusCode == nil && bifrostErr.Error != nil && bifrostErr.Error.Type == nil && + bifrostErr.Error.Message == "" && bifrostErr.Error.Error == nil { + return resp, nil + } + return resp, bifrostErr + } + return resp, nil +} + // resetPluginPipeline resets a PluginPipeline instance for reuse. // IMPORTANT: drainAndAttachPluginLogs must be called on the root BifrostContext // BEFORE this method, because it calls ReleasePluginScope on cached scoped contexts @@ -7129,25 +6999,6 @@ func (bifrost *Bifrost) releaseBifrostRequest(req *schemas.BifrostRequest) { bifrost.bifrostRequestPool.Put(req) } -// resetMCPRequest resets a BifrostMCPRequest instance for reuse -func resetMCPRequest(req *schemas.BifrostMCPRequest) { - req.RequestType = "" - req.ChatAssistantMessageToolCall = nil - req.ResponsesToolMessage = nil -} - -// getMCPRequest gets a BifrostMCPRequest from the pool -func (bifrost *Bifrost) getMCPRequest() *schemas.BifrostMCPRequest { - req := bifrost.mcpRequestPool.Get().(*schemas.BifrostMCPRequest) - return req -} - -// releaseMCPRequest returns a BifrostMCPRequest to the pool -func (bifrost *Bifrost) releaseMCPRequest(req *schemas.BifrostMCPRequest) { - resetMCPRequest(req) - bifrost.mcpRequestPool.Put(req) -} - // getAllSupportedKeys retrieves all valid keys for a ListModels request. // allowing the provider to aggregate results from multiple keys. func (bifrost *Bifrost) getAllSupportedKeys(ctx *schemas.BifrostContext, providerKey schemas.ModelProvider, baseProviderType schemas.ModelProvider) ([]schemas.Key, error) { diff --git a/core/internal/mcptests/agent_adapter_test.go b/core/internal/mcptests/agent_adapter_test.go index fcf97b4668f..040ecf1c218 100644 --- a/core/internal/mcptests/agent_adapter_test.go +++ b/core/internal/mcptests/agent_adapter_test.go @@ -63,9 +63,6 @@ func TestAgent_Adapter_ResponsesFormat_BasicLoop(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForResponsesRequest( ctx, req, initialResponse, mocker.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -113,9 +110,6 @@ func TestAgent_Adapter_ResponsesFormat_EmptyToolResult(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForResponsesRequest( ctx, req, initialResponse, mocker.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -185,9 +179,6 @@ func TestAgent_Adapter_ResponsesFormat_MultipleToolCalls(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForResponsesRequest( ctx, req, initialResponse, mocker.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -237,9 +228,6 @@ func TestAgent_Adapter_ResponsesFormat_MixedPermissions(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForResponsesRequest( ctx, req, initialResponse, mocker.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -302,9 +290,6 @@ func TestAgent_Adapter_ResponsesFormat_STDIO(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForResponsesRequest( ctx, req, initialResponse, mocker.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -368,9 +353,6 @@ func TestAgent_Adapter_ResponsesFormat_DeepChain(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForResponsesRequest( ctx, req, initialResponse, mocker.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -422,9 +404,6 @@ func TestAgent_Adapter_ResponsesFormat_ErrorHandling(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForResponsesRequest( ctx, req, initialResponse, mocker.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -502,9 +481,6 @@ func TestAgent_Adapter_ChatAndResponsesParity(t *testing.T) { chatResult, chatErr := managerChat.CheckAndExecuteAgentForChatRequest( ctxChat, chatReq, chatInitialResponse, mockerChat.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return managerChat.ExecuteToolCall(ctx, request) - }, ) // Execute Responses API @@ -519,9 +495,6 @@ func TestAgent_Adapter_ChatAndResponsesParity(t *testing.T) { responsesResult, responsesErr := managerResponses.CheckAndExecuteAgentForResponsesRequest( ctxResponses, responsesReq, responsesInitialResponse, mockerResponses.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return managerResponses.ExecuteToolCall(ctx, request) - }, ) // Assertions: Both should complete successfully diff --git a/core/internal/mcptests/agent_basic_test.go b/core/internal/mcptests/agent_basic_test.go index 7d16ca5e49d..4210ac49ca3 100644 --- a/core/internal/mcptests/agent_basic_test.go +++ b/core/internal/mcptests/agent_basic_test.go @@ -107,10 +107,6 @@ func TestAgent_BasicLoop(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - // Use real tool execution - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr, "agent loop should complete successfully") @@ -170,9 +166,6 @@ func TestAgent_BasicLoop_ChatFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -233,9 +226,6 @@ func TestAgent_BasicLoop_ResponsesFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -294,9 +284,6 @@ func TestAgent_SingleIteration(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -372,9 +359,6 @@ func TestAgent_MultipleIterations(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -429,9 +413,6 @@ func TestAgent_NoToolCalls(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -496,9 +477,6 @@ func TestAgent_MixedAutoAndNonAutoTools(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -575,9 +553,6 @@ func TestAgent_OnlyAutoTools(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -637,9 +612,6 @@ func TestAgent_OnlyNonAutoTools(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) diff --git a/core/internal/mcptests/agent_context_filtering_test.go b/core/internal/mcptests/agent_context_filtering_test.go index 8803d9730ed..79b16908bd7 100644 --- a/core/internal/mcptests/agent_context_filtering_test.go +++ b/core/internal/mcptests/agent_context_filtering_test.go @@ -70,9 +70,6 @@ func TestAgent_ContextToolFilter_Whitelist(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -117,9 +114,6 @@ func TestAgent_ContextToolFilter_BlockedToolError(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -175,9 +169,6 @@ func TestAgent_ContextClientFilter_Whitelist(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -226,9 +217,6 @@ func TestAgent_ContextNarrowing_AutoExecute(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -275,9 +263,6 @@ func TestAgent_ContextToolFilter_EmptyList(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -340,9 +325,6 @@ func TestAgent_ContextToolFilter_WildcardOverride(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -412,9 +394,6 @@ func TestAgent_ContextClientFilter_MultipleClients(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -465,9 +444,6 @@ func TestAgent_ContextToolFilter_ParallelMixed(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions diff --git a/core/internal/mcptests/agent_error_handling_test.go b/core/internal/mcptests/agent_error_handling_test.go index 7e1c016a038..eebaf1fa807 100644 --- a/core/internal/mcptests/agent_error_handling_test.go +++ b/core/internal/mcptests/agent_error_handling_test.go @@ -91,9 +91,6 @@ func TestAgent_ErrorHandling_AllToolsFail(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Agent should handle all failures gracefully @@ -173,9 +170,6 @@ func TestAgent_ErrorHandling_TimeoutInLoop(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Should handle timeout gracefully @@ -340,9 +334,6 @@ func TestAgent_ErrorHandling_PartialBatchFailure(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr, "partial failures should not crash agent") @@ -433,9 +424,6 @@ func TestAgent_ErrorHandling_RecoveryAndContinuation(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr, "agent should recover from error") @@ -575,9 +563,6 @@ func TestAgent_ErrorHandling_MultipleErrorsInSequence(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr, "agent should handle multiple sequential errors") @@ -649,9 +634,6 @@ func TestAgent_ErrorHandling_ErrorMessagePreservation(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) diff --git a/core/internal/mcptests/agent_filtering_test.go b/core/internal/mcptests/agent_filtering_test.go index 1058f5a7f25..e1a7f1cd64c 100644 --- a/core/internal/mcptests/agent_filtering_test.go +++ b/core/internal/mcptests/agent_filtering_test.go @@ -56,9 +56,6 @@ func TestAgent_ToolAllowedNotAutoExecute(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -118,9 +115,6 @@ func TestAgent_ToolAllowedAndAutoExecute(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -173,9 +167,6 @@ func TestAgent_ToolNotAllowed(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -234,9 +225,6 @@ func TestAgent_ToolNotInAutoExecuteList(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -306,9 +294,6 @@ func TestAgent_ComplexFiltering_Scenario1(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -373,9 +358,6 @@ func TestAgent_ComplexFiltering_Scenario2(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -439,9 +421,6 @@ func TestAgent_ComplexFiltering_Scenario3(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -496,9 +475,6 @@ func TestAgent_ComplexFiltering_ContextOverride(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -581,9 +557,6 @@ func TestAgent_FilteringWithMultipleClients(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -668,9 +641,6 @@ func TestAgent_ToolConflictInAgentMode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -757,9 +727,6 @@ func TestAgent_AllAutoExecuteScenarios(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) if scenario.ShouldAutoExecute { @@ -830,9 +797,6 @@ func TestAgent_Filtering_ChatFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -889,9 +853,6 @@ func TestAgent_Filtering_ResponsesFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) diff --git a/core/internal/mcptests/agent_limits_test.go b/core/internal/mcptests/agent_limits_test.go index 3c4d21ec907..f4fce6c863f 100644 --- a/core/internal/mcptests/agent_limits_test.go +++ b/core/internal/mcptests/agent_limits_test.go @@ -76,9 +76,6 @@ func TestAgent_MaxDepthEnforcement(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -143,9 +140,6 @@ func TestAgent_MaxDepthCustomValue(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -206,9 +200,6 @@ func TestAgent_MaxDepthReached_ChatFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -273,9 +264,6 @@ func TestAgent_MaxDepthReached_ResponsesFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -364,9 +352,6 @@ func TestAgent_MaxDepth_CodeMode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -439,9 +424,6 @@ func TestAgent_MaxDepth_CodeMode_ChatFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -508,9 +490,6 @@ func TestAgent_MaxDepth_CodeMode_ResponsesFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -589,9 +568,6 @@ func TestAgent_Timeout(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // MUST have timeout error @@ -654,9 +630,6 @@ func TestAgent_TimeoutDuringExecution(t *testing.T) { }, mockLLM.chatResponses[0], mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // MUST timeout during execution @@ -703,10 +676,7 @@ func TestAgent_Timeout_ChatFormat(t *testing.T) { _, bifrostErr := manager.CheckAndExecuteAgentForChatRequest(ctx, &schemas.BifrostChatRequest{Provider: schemas.OpenAI, Model: "gpt-4o", Input: []schemas.ChatMessage{{Role: schemas.ChatMessageRoleUser, Content: &schemas.ChatMessageContent{ContentStr: schemas.Ptr("Test")}}}}, - mockLLM.chatResponses[0], mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }) + mockLLM.chatResponses[0], mockLLM.MakeChatRequest) require.NotNil(t, bifrostErr, "Chat format timeout not enforced!") t.Logf("✅ Chat format timeout enforced: %v", bifrostErr.Error) @@ -752,10 +722,7 @@ func TestAgent_Timeout_ResponsesFormat(t *testing.T) { &schemas.BifrostResponsesRequest{Provider: schemas.OpenAI, Model: "gpt-4o", Input: []schemas.ResponsesMessage{{Type: schemas.Ptr(schemas.ResponsesMessageTypeMessage), Role: schemas.Ptr(schemas.ResponsesInputMessageRoleUser), Content: &schemas.ResponsesMessageContent{ContentStr: schemas.Ptr("Test")}}}}, - mockLLM.responsesResponses[0], mockLLM.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }) + mockLLM.responsesResponses[0], mockLLM.MakeResponsesRequest) require.NotNil(t, bifrostErr, "Responses format timeout not enforced!") t.Logf("✅ Responses format timeout enforced: %v", bifrostErr.Error) @@ -814,9 +781,6 @@ func TestAgent_ErrorPropagation(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Error should be propagated, or tool result should contain error @@ -881,9 +845,6 @@ func TestAgent_ErrorInMiddleOfLoop(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // First tool should have executed successfully @@ -936,9 +897,6 @@ func TestAgent_LLMError(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // LLM error should be returned @@ -1008,9 +966,6 @@ func TestAgent_MaxDepthAndTimeout(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Whichever limit hits first should stop the agent @@ -1073,9 +1028,6 @@ func TestAgent_MaxDepthZero(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Should return immediately with tool calls @@ -1133,9 +1085,6 @@ func TestAgent_ParallelToolExecution(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -1198,9 +1147,6 @@ func TestAgent_IterationTracking(t *testing.T) { originalReq, initialResponse, trackingMockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) diff --git a/core/internal/mcptests/agent_mixed_permissions_test.go b/core/internal/mcptests/agent_mixed_permissions_test.go index 3cd6213cb50..32d058179bd 100644 --- a/core/internal/mcptests/agent_mixed_permissions_test.go +++ b/core/internal/mcptests/agent_mixed_permissions_test.go @@ -66,9 +66,6 @@ func TestAgent_MixedPermissions_ThreeClients(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -144,9 +141,6 @@ func TestAgent_MixedPermissions_AllBlocked(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -203,9 +197,6 @@ func TestAgent_MixedPermissions_WildcardAutoExecute(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -270,9 +261,6 @@ func TestAgent_MixedPermissions_PartialExecution(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -366,9 +354,6 @@ func TestAgent_MixedPermissions_MultipleSTDIOSamePermissions(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -427,9 +412,6 @@ func TestAgent_MixedPermissions_ContextFilteringOverride(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -497,9 +479,6 @@ func TestAgent_MixedPermissions_SpecificToolNames(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -576,9 +555,6 @@ func TestAgent_MixedPermissions_AllAutoExecute(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions diff --git a/core/internal/mcptests/agent_multiconnection_test.go b/core/internal/mcptests/agent_multiconnection_test.go index 7b5892d7035..1b41130e73f 100644 --- a/core/internal/mcptests/agent_multiconnection_test.go +++ b/core/internal/mcptests/agent_multiconnection_test.go @@ -65,9 +65,6 @@ func TestAgent_MultiConnection_AllTypes(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -132,9 +129,6 @@ func TestAgent_MultiConnection_MixedPermissions(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -221,9 +215,6 @@ func TestAgent_MultiConnection_SequentialAfterParallel(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -274,9 +265,6 @@ func TestAgent_MultiConnection_ErrorInSTDIO(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -362,9 +350,6 @@ func TestAgent_MultiConnection_LargeParallelBatch(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions diff --git a/core/internal/mcptests/agent_parallel_execution_test.go b/core/internal/mcptests/agent_parallel_execution_test.go index 47ea6884ec1..b308b3694a4 100644 --- a/core/internal/mcptests/agent_parallel_execution_test.go +++ b/core/internal/mcptests/agent_parallel_execution_test.go @@ -91,9 +91,6 @@ func TestAgent_ParallelExecution_ResultOrdering(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr, "agent execution should succeed") @@ -181,9 +178,6 @@ func TestAgent_ParallelExecution_PartialFailures(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr, "agent should handle partial failures gracefully") @@ -275,9 +269,6 @@ func TestAgent_ParallelExecution_RaceConditions(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -368,9 +359,6 @@ func TestAgent_ParallelExecution_LargeBatch(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) elapsed := time.Since(start) @@ -465,9 +453,6 @@ func TestAgent_ParallelExecution_MixedOutcomes(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Should complete despite mixed outcomes @@ -560,9 +545,6 @@ func TestAgent_ParallelExecution_ResultCollectionOrder(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) diff --git a/core/internal/mcptests/agent_request_id_test.go b/core/internal/mcptests/agent_request_id_test.go index 530d2e0d8d4..b96254a78e4 100644 --- a/core/internal/mcptests/agent_request_id_test.go +++ b/core/internal/mcptests/agent_request_id_test.go @@ -144,9 +144,6 @@ func TestAgent_RequestID_Propagation(t *testing.T) { req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -236,9 +233,6 @@ func TestAgent_RequestID_PreservationAcrossDepth(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -311,9 +305,6 @@ func TestAgent_RequestID_NoGeneratorFunction(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -380,9 +371,6 @@ func TestAgent_RequestID_EmptyGeneratorResult(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -452,9 +440,6 @@ func TestAgent_RequestID_SequentialUpdates(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -526,9 +511,6 @@ func TestAgent_RequestID_MixedAutoAndNonAuto(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions diff --git a/core/internal/mcptests/agent_state_transitions_test.go b/core/internal/mcptests/agent_state_transitions_test.go index f3c6fd27d6b..a73b2fa3766 100644 --- a/core/internal/mcptests/agent_state_transitions_test.go +++ b/core/internal/mcptests/agent_state_transitions_test.go @@ -125,9 +125,6 @@ func TestAgent_StateTransition_LargeMixedToolBatch(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr, "should handle large mixed batch") @@ -197,9 +194,6 @@ func TestAgent_StateTransition_DepthCountingBasic(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -283,9 +277,6 @@ func TestAgent_StateTransition_AlternatingAutoNonAuto(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -367,9 +358,6 @@ func TestAgent_StateTransition_EmptyToolCallsList(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -459,9 +447,6 @@ func TestAgent_StateTransition_AllToolsFilteredOut(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -561,9 +546,6 @@ func TestAgent_StateTransition_StateConsistency(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -635,9 +617,6 @@ func TestAgent_StateTransition_BoundaryConditions(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) diff --git a/core/internal/mcptests/agent_test_helpers.go b/core/internal/mcptests/agent_test_helpers.go index 85512dcce6f..674f4bd2ea2 100644 --- a/core/internal/mcptests/agent_test_helpers.go +++ b/core/internal/mcptests/agent_test_helpers.go @@ -624,9 +624,6 @@ func RunAgentScenario(t *testing.T, scenario AgentScenario) { req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Run assertions @@ -678,9 +675,6 @@ func SimpleAgentTest(t *testing.T, name string, config AgentTestConfig, response req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) assertions(t, response, bifrostErr, mocker) diff --git a/core/internal/mcptests/agent_test_helpers_example_test.go b/core/internal/mcptests/agent_test_helpers_example_test.go index 8e840fda967..b42ebc450f3 100644 --- a/core/internal/mcptests/agent_test_helpers_example_test.go +++ b/core/internal/mcptests/agent_test_helpers_example_test.go @@ -53,9 +53,6 @@ func TestAgentHelpers_Example_SimpleInProcessAgent(t *testing.T) { req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions using agent-specific helpers @@ -101,9 +98,6 @@ func TestAgentHelpers_Example_MultiConnectionTypes(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assert parallel execution @@ -145,9 +139,6 @@ func TestAgentHelpers_Example_ContextFiltering(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assert echo executed but weather blocked (agent stops with error at turn 2) @@ -210,9 +201,6 @@ func TestAgentHelpers_Example_MaxDepthLimit(t *testing.T) { result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest( ctx, req, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Agent should stop at max depth (initial call + up to maxDepth-1 continuations) diff --git a/core/internal/mcptests/codemode_agent_multiturn_test.go b/core/internal/mcptests/codemode_agent_multiturn_test.go index 419ef437879..40fbe1765f8 100644 --- a/core/internal/mcptests/codemode_agent_multiturn_test.go +++ b/core/internal/mcptests/codemode_agent_multiturn_test.go @@ -135,9 +135,6 @@ func TestCodeMode_Agent_MultiTurn_CodeChaining(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -259,9 +256,6 @@ else: originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, agentErr) @@ -357,9 +351,6 @@ else: originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, agentErr) @@ -468,9 +459,6 @@ func TestCodeMode_Agent_MultiTurn_ContextFilterOverride(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, agentErr) @@ -578,9 +566,6 @@ func TestCodeMode_Agent_MultiTurn_MaxDepth(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, agentErr) @@ -689,9 +674,6 @@ func TestCodeMode_Agent_MultiTurn_ErrorRecovery(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, agentErr) diff --git a/core/internal/mcptests/codemode_agent_singleturn_test.go b/core/internal/mcptests/codemode_agent_singleturn_test.go index dab012d572b..ec84d28caaf 100644 --- a/core/internal/mcptests/codemode_agent_singleturn_test.go +++ b/core/internal/mcptests/codemode_agent_singleturn_test.go @@ -119,9 +119,6 @@ func TestCodeMode_Agent_AutoExecuteSingleTool(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions @@ -225,9 +222,6 @@ func TestCodeMode_Agent_NonAutoToolInCode(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Assertions diff --git a/core/internal/mcptests/codemode_agent_test.go b/core/internal/mcptests/codemode_agent_test.go index 668482de5ac..1098b5fe98d 100644 --- a/core/internal/mcptests/codemode_agent_test.go +++ b/core/internal/mcptests/codemode_agent_test.go @@ -66,9 +66,6 @@ func TestCodeModeAgent_Basic(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err, "agent loop should complete successfully") @@ -136,9 +133,6 @@ func TestCodeModeAgent_NonAutoToolInCode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -212,9 +206,6 @@ func TestCodeModeAgent_AutoToolInCode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -280,9 +271,6 @@ func TestCodeModeAgent_MixedToolsInCode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -357,9 +345,6 @@ func TestCodeModeAgent_NoToolCallsInCode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -426,9 +411,6 @@ func TestCodeModeAgent_FilteringInCode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -492,9 +474,6 @@ func TestCodeModeAgent_AutoExecuteFiltering(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -575,9 +554,6 @@ func TestCodeModeAgent_MaxDepth(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -645,9 +621,6 @@ func TestCodeModeAgent_MaxDepth_ChatFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -718,9 +691,6 @@ func TestCodeModeAgent_MaxDepth_ResponsesFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -787,9 +757,6 @@ func TestCodeModeAgent_Timeout(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -850,9 +817,6 @@ func TestCodeModeAgent_Timeout_ChatFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -913,9 +877,6 @@ func TestCodeModeAgent_Timeout_ResponsesFormat(t *testing.T) { originalReq, initialResponse, mockLLM.MakeResponsesRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -975,9 +936,6 @@ func TestCodeModeAgent_ErrorInCode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) @@ -1040,9 +998,6 @@ func TestCodeModeAgent_ToolErrorInCode(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, err) diff --git a/core/internal/mcptests/codemode_vs_noncodemode_test.go b/core/internal/mcptests/codemode_vs_noncodemode_test.go index 61ad852ab82..4c75ac95e33 100644 --- a/core/internal/mcptests/codemode_vs_noncodemode_test.go +++ b/core/internal/mcptests/codemode_vs_noncodemode_test.go @@ -341,9 +341,6 @@ func TestCodeMode_Agent_MixedCodeModeWithApproval(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, agentErr) @@ -449,9 +446,6 @@ func TestCodeMode_Agent_CodeModeInCode_NonCodeModeDirect(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, agentErr) @@ -555,9 +549,6 @@ func TestCodeMode_Agent_PartialApprovalMixed(t *testing.T) { originalReq, initialResponse, mocker.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, agentErr) diff --git a/core/internal/mcptests/connect_ping_listtools_test.go b/core/internal/mcptests/connect_ping_listtools_test.go new file mode 100644 index 00000000000..7d7f8b5455c --- /dev/null +++ b/core/internal/mcptests/connect_ping_listtools_test.go @@ -0,0 +1,599 @@ +package mcptests + +import ( + "context" + "fmt" + "testing" + "time" + + core "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/mcp" + "github.com/maximhq/bifrost/core/schemas" + mcpgo "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ============================================================================= +// SHARED HELPERS +// ============================================================================= + +// buildInProcessServer creates a fresh mcp-go server with a deterministic set of +// tools and returns it. The server is independent of the bifrost-internal server +// so AddClient flows go through the full connect/list_tools gate. +func buildInProcessServer(t *testing.T) *server.MCPServer { + t.Helper() + s := server.NewMCPServer("test-inproc", "1.2.3", server.WithToolCapabilities(true)) + + // Tool A — echo + echoTool := mcpgo.NewTool("echo", + mcpgo.WithDescription("Echo tool"), + mcpgo.WithString("message", mcpgo.Required(), mcpgo.Description("message")), + ) + s.AddTool(echoTool, func(ctx context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { + msg, _ := req.GetArguments()["message"].(string) + return mcpgo.NewToolResultText(msg), nil + }) + + // Tool B — adder (separate tool so PostHook filtering tests can drop one and keep the other) + addTool := mcpgo.NewTool("add", + mcpgo.WithDescription("Adds two numbers"), + mcpgo.WithNumber("x", mcpgo.Required(), mcpgo.Description("x")), + mcpgo.WithNumber("y", mcpgo.Required(), mcpgo.Description("y")), + ) + s.AddTool(addTool, func(ctx context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { + args := req.GetArguments() + x, _ := args["x"].(float64) + y, _ := args["y"].(float64) + return mcpgo.NewToolResultText(fmt.Sprintf("%v", x+y)), nil + }) + + return s +} + +// inProcessClientConfig builds a Client config wrapping the given server, ready +// for AddClient. ID embeds clientName for easy identification in test assertions. +func inProcessClientConfig(clientName string, s *server.MCPServer) *schemas.MCPClientConfig { + return &schemas.MCPClientConfig{ + ID: clientName + "-id", + Name: clientName, + ConnectionType: schemas.MCPConnectionTypeInProcess, + InProcessServer: s, + ToolsToExecute: []string{"*"}, + } +} + +// setupBifrostWithPlugins returns a manager + bifrost where the manager's plugin +// pipeline is wired to a Bifrost instance carrying the given MCP plugins. Plugins +// fire for any AddClient performed after this returns. +func setupBifrostWithPlugins(t *testing.T, plugins []schemas.MCPPlugin) (*mcp.MCPManager, *core.Bifrost) { + t.Helper() + manager := setupMCPManager(t) + bf, err := core.Init(context.Background(), schemas.BifrostConfig{ + Account: &testAccount{}, + MCPPlugins: plugins, + Logger: core.NewDefaultLogger(schemas.LogLevelError), + }) + require.NoError(t, err) + bf.SetMCPManager(manager) + return manager, bf +} + +// ============================================================================= +// CONNECT HOOK TESTS +// ============================================================================= + +func TestConnectHook_FiresOnAddClient(t *testing.T) { + t.Parallel() + + plugin := NewTestConnectPlugin() + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + cfg := inProcessClientConfig("connect_fires", buildInProcessServer(t)) + require.NoError(t, manager.AddClient(cfg)) + + pre := plugin.GetPreHookCalls() + post := plugin.GetPostHookCalls() + require.Len(t, pre, 1, "Connect PreHook should fire exactly once for AddClient") + require.Len(t, post, 1, "Connect PostHook should fire exactly once for AddClient") + + // Verify PreHook saw the right typed sub-request. + req := pre[0].ConnectRequest + require.NotNil(t, req, "typed Connect captures land in ConnectRequest, not the envelope Request field") + assert.Equal(t, "connect_fires", req.ClientName) + assert.Equal(t, schemas.MCPConnectionTypeInProcess, req.ConnectionType) +} + +func TestConnectHook_PostHookPopulatesServerInfo(t *testing.T) { + t.Parallel() + + plugin := NewTestConnectPlugin() + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + require.NoError(t, manager.AddClient(inProcessClientConfig("server_info", buildInProcessServer(t)))) + + post := plugin.GetPostHookCalls() + require.Len(t, post, 1) + resp := post[0].ConnectResponse + require.NotNil(t, resp, "typed Connect captures land in ConnectResponse") + + si := resp.ServerInfo + require.NotNil(t, si, "ServerInfo must be populated from initialize handshake") + assert.Equal(t, "test-inproc", si.Name) + assert.Equal(t, "1.2.3", si.Version) + assert.NotEmpty(t, resp.ProtocolVersion) + require.NotNil(t, resp.ServerCapabilities) + assert.True(t, resp.ServerCapabilities.Tools, "server advertises tool capability") + + // ExtraFields on the typed sub-response carries Latency + ClientName backfill. + assert.Greater(t, resp.ExtraFields.Latency, int64(-1), "Latency should be non-negative") + assert.Equal(t, "server_info", resp.ExtraFields.ClientName, "ClientName backfilled via PopulateExtraFields") + + // Typed plugins also see ClientName on the captured sub-request. + pre := plugin.GetPreHookCalls() + require.NotEmpty(t, pre) + require.NotNil(t, pre[0].ConnectRequest) + assert.Equal(t, "server_info", pre[0].ConnectRequest.ClientName) +} + +func TestConnectHook_PreHookShortCircuitError_FailsAddClient(t *testing.T) { + t.Parallel() + + plugin := NewTestConnectPlugin() + plugin.SetShortCircuitError(&schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "blocked by governance"}, + }) + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + err := manager.AddClient(inProcessClientConfig("blocked_client", buildInProcessServer(t))) + require.Error(t, err, "Plugin error short-circuit should fail AddClient") + assert.Contains(t, err.Error(), "blocked by governance") + + // PostHook should still have fired for the executed PreHook plugins. + assert.Len(t, plugin.GetPostHookCalls(), 0, "PostHook only fires on success or recovery; raw error short-circuit yields no response") + assert.Len(t, plugin.GetPreHookCalls(), 1, "PreHook ran once before short-circuit") + + // Verify no client was registered. + clients := manager.GetClients() + for _, c := range clients { + assert.NotEqual(t, "blocked_client", c.Name, "no client should be registered after error short-circuit") + } +} + +func TestConnectHook_PreHookShortCircuitResponse_RegistersWithEmptyTools(t *testing.T) { + t.Parallel() + + plugin := NewTestConnectPlugin() + plugin.SetShortCircuitResponse(&schemas.BifrostMCPConnectResponse{ + ServerInfo: &schemas.MCPServerInfo{Name: "synthetic_client", Version: "0.0.0"}, + ConnectionInfo: &schemas.MCPClientConnectionInfo{ + Type: schemas.MCPConnectionTypeInProcess, + }, + }) + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + // AddClient should succeed (no wire dial happens) — documented Connect-success + // short-circuit gotcha: client registered as connected with no live transport. + require.NoError(t, manager.AddClient(inProcessClientConfig("synthetic_client", buildInProcessServer(t)))) + + clients := manager.GetClients() + var found *schemas.MCPClientState + for i := range clients { + if clients[i].Name == "synthetic_client" { + found = &clients[i] + break + } + } + require.NotNil(t, found, "client should be registered even with synthetic connect") + assert.Empty(t, found.ToolMap, "synthetic-connect client has no tools (list_tools is never called)") +} + +func TestConnectHook_AuthorizationHidden_HeadersAuth(t *testing.T) { + t.Parallel() + + // SECURITY: PreHook plugins must NOT see the Authorization header. The connect + // gate strips Authorization from headers exposed to plugins and re-injects it + // only after all PreHooks have run. Test via MCPAuthTypeHeaders so we don't need + // a live OAuth provider. + // + // We short-circuit after capture so the test doesn't wait on the unreachable + // transport retry loop. The strip happens before any PreHook runs, so capturing + // the request once is sufficient to prove it. + plugin := NewTestConnectPlugin() + plugin.SetShortCircuitError(&schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "captured, aborting"}, + }) + + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + url := *schemas.NewEnvVar("http://example.invalid") + cfg := &schemas.MCPClientConfig{ + ID: "auth_strip-id", + Name: "auth_strip", + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: &url, + AuthType: schemas.MCPAuthTypeHeaders, + Headers: map[string]schemas.EnvVar{ + "Authorization": *schemas.NewEnvVar("Bearer super-secret-token"), + "X-Custom": *schemas.NewEnvVar("plugin-visible"), + }, + ToolsToExecute: []string{"*"}, + } + // Short-circuit returns an error → AddClient fails. That's expected. + _ = manager.AddClient(cfg) + + calls := plugin.GetPreHookCalls() + require.NotEmpty(t, calls, "PreHook should have fired before short-circuit") + + req := calls[0].ConnectRequest + require.NotNil(t, req) + require.NotNil(t, req.Headers, "Headers should be populated (we set X-Custom)") + + // Authorization must NOT be present in the headers plugins see. + _, hasAuth := req.Headers["Authorization"] + assert.False(t, hasAuth, "Authorization header must be stripped before PreHooks run") + // Verify the secret never appeared anywhere in the visible header values. + for k, v := range req.Headers { + assert.NotContains(t, v, "super-secret-token", + "bearer token should not leak via any header (%q)", k) + } + // Other user-configured headers should still be visible. + assert.Equal(t, "plugin-visible", req.Headers["X-Custom"], + "Non-Authorization headers should remain visible to plugins") +} + +func TestConnectHook_OnlyFiresForConnectRequestType(t *testing.T) { + t.Parallel() + + // TestConnectPlugin filters by RequestType. Verify it ignores execute-tool flows. + connectPlugin := NewTestConnectPlugin() + logPlugin := NewTestLoggingPlugin() // captures everything for cross-reference + manager, bf := setupBifrostWithPlugins(t, []schemas.MCPPlugin{connectPlugin, logPlugin}) + + // Trigger execute-tool path (no connect involved on the bifrost-internal client). + require.NoError(t, RegisterEchoTool(manager)) + echoCall := GetSampleEchoToolCall("filter_test", "hello") + _, bifrostErr := bf.ExecuteChatMCPTool(createTestContext(), &echoCall) + require.Nil(t, bifrostErr) + + // Connect plugin must not have fired. + assert.Empty(t, connectPlugin.GetPreHookCalls(), "Connect plugin should ignore execute-tool requests") + assert.Empty(t, connectPlugin.GetPostHookCalls(), "Connect plugin should ignore execute-tool responses") + // Logging plugin captures the execute-tool flow (as a cross-check that the test ran). + assert.GreaterOrEqual(t, logPlugin.GetPreHookCallCount(), 1) +} + +// ============================================================================= +// LISTTOOLS HOOK TESTS +// ============================================================================= + +func TestListToolsHook_FiresOnAddClient(t *testing.T) { + t.Parallel() + + plugin := NewTestListToolsPlugin() + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + require.NoError(t, manager.AddClient(inProcessClientConfig("list_fires", buildInProcessServer(t)))) + + pre := plugin.GetPreHookCalls() + post := plugin.GetPostHookCalls() + require.Len(t, pre, 1, "ListTools PreHook should fire once during AddClient (post-init tool retrieval)") + require.Len(t, post, 1, "ListTools PostHook should fire once") + + // Verify response shape. + resp := post[0].Response + require.NotNil(t, resp) + require.NotNil(t, resp.BifrostMCPListToolsResponse) + assert.Equal(t, 2, resp.RawToolCount, "raw count should reflect both tools (echo + add)") + assert.Len(t, resp.Tools, 2, "filtered tools should match (no name violations in this set)") + // Tools are prefixed with client name. + _, hasEcho := resp.Tools["list_fires-echo"] + _, hasAdd := resp.Tools["list_fires-add"] + assert.True(t, hasEcho, "echo tool should be present with client prefix") + assert.True(t, hasAdd, "add tool should be present with client prefix") +} + +func TestListToolsHook_PostHookFilterAppliedToClientState(t *testing.T) { + t.Parallel() + + plugin := NewTestListToolsPlugin() + // Drop the "add" tool; keep "echo". + plugin.SetPostHookFilter(func(prefixedName string) bool { + return prefixedName == "list_filter-echo" + }) + + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + require.NoError(t, manager.AddClient(inProcessClientConfig("list_filter", buildInProcessServer(t)))) + + // Verify the filtered set landed in the manager's stored ToolMap (not just the + // gate response). The connect path applies the gate result to clientState.ToolMap. + clients := manager.GetClients() + var target *schemas.MCPClientState + for i := range clients { + if clients[i].Name == "list_filter" { + target = &clients[i] + break + } + } + require.NotNil(t, target) + assert.Len(t, target.ToolMap, 1, "PostHook filter should have removed 'add' tool") + _, hasEcho := target.ToolMap["list_filter-echo"] + _, hasAdd := target.ToolMap["list_filter-add"] + assert.True(t, hasEcho, "filtered ToolMap should keep echo") + assert.False(t, hasAdd, "filtered ToolMap should drop add") +} + +func TestListToolsHook_PreHookShortCircuitWithSyntheticTools(t *testing.T) { + t.Parallel() + + plugin := NewTestListToolsPlugin() + synthetic := &schemas.BifrostMCPListToolsResponse{ + Tools: map[string]schemas.ChatTool{ + "synthetic-tool": { + Type: schemas.ChatToolTypeFunction, + Function: &schemas.ChatToolFunction{ + Name: "synthetic-tool", + Description: schemas.Ptr("Plugin-injected tool"), + }, + }, + }, + ToolNameMapping: map[string]string{"synthetic_tool": "synthetic-tool"}, + } + plugin.SetShortCircuitResponse(synthetic) + + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + require.NoError(t, manager.AddClient(inProcessClientConfig("list_synth", buildInProcessServer(t)))) + + clients := manager.GetClients() + var target *schemas.MCPClientState + for i := range clients { + if clients[i].Name == "list_synth" { + target = &clients[i] + break + } + } + require.NotNil(t, target) + // The synthetic tool list should have replaced the real server's tools. + _, hasSynthetic := target.ToolMap["synthetic-tool"] + _, hasEcho := target.ToolMap["list_synth-echo"] + assert.True(t, hasSynthetic, "synthetic tool should be in ToolMap from short-circuit") + assert.False(t, hasEcho, "real server tools should not appear when PreHook short-circuited") +} + +func TestListToolsHook_PreHookShortCircuitError_LeavesEmptyToolMap(t *testing.T) { + t.Parallel() + + plugin := NewTestListToolsPlugin() + plugin.SetShortCircuitError(&schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "list_tools blocked"}, + }) + + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + // AddClient should still succeed — the connect path tolerates list_tools failure + // and falls back to empty tools (matching pre-plugin behavior). + require.NoError(t, manager.AddClient(inProcessClientConfig("list_err", buildInProcessServer(t)))) + + clients := manager.GetClients() + var target *schemas.MCPClientState + for i := range clients { + if clients[i].Name == "list_err" { + target = &clients[i] + break + } + } + require.NotNil(t, target) + assert.Empty(t, target.ToolMap, "list_tools error short-circuit should result in empty ToolMap") +} + +func TestListToolsHook_FiresOnConnectAndAgain(t *testing.T) { + t.Parallel() + + // Verifies that re-establishing a connection re-fires the list_tools gate. + plugin := NewTestListToolsPlugin() + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + cfg := inProcessClientConfig("list_reconnect", buildInProcessServer(t)) + require.NoError(t, manager.AddClient(cfg)) + require.Len(t, plugin.GetPreHookCalls(), 1, "first AddClient should fire list_tools once") + + // Reconnect: this tears down and re-establishes the client, firing list_tools again. + require.NoError(t, manager.ReconnectClient(cfg.ID)) + require.GreaterOrEqual(t, len(plugin.GetPreHookCalls()), 2, "ReconnectClient should re-fire list_tools") +} + +// ============================================================================= +// PING HOOK TESTS +// ============================================================================= + +func TestPingHook_FiresViaHealthMonitor(t *testing.T) { + t.Parallel() + + plugin := NewTestPingPlugin() + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + cfg := inProcessClientConfig("ping_fires", buildInProcessServer(t)) + require.NoError(t, manager.AddClient(cfg)) + + // AddClient starts its own health monitor at 10s interval — far too slow for + // tests. Spin up a dedicated monitor at 10ms instead. + monitor := mcp.NewClientHealthMonitor(manager, cfg.ID, 10*time.Millisecond, true, core.NewDefaultLogger(schemas.LogLevelError)) + monitor.Start() + defer monitor.Stop() + + // Wait for at least a couple of complete tick cycles (both PreHook and PostHook). + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if plugin.GetPreHookCallCount() >= 2 && plugin.GetPostHookCallCount() >= 2 { + break + } + time.Sleep(10 * time.Millisecond) + } + + preCount := plugin.GetPreHookCallCount() + postCount := plugin.GetPostHookCallCount() + require.GreaterOrEqual(t, preCount, 2, "Ping PreHook should fire on each health-check tick") + require.GreaterOrEqual(t, postCount, 2, "Ping PostHook should fire on each successful ping") + + // Verify request shape. + pre := plugin.GetPreHookCalls() + for _, e := range pre { + assert.Equal(t, schemas.MCPRequestTypePing, e.Request.RequestType) + assert.Equal(t, "ping_fires", e.Request.ClientName) + } +} + +func TestPingHook_PreHookShortCircuitHealthy(t *testing.T) { + t.Parallel() + + plugin := NewTestPingPlugin() + plugin.SetShortCircuitHealthy(true) + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + + cfg := inProcessClientConfig("ping_healthy", buildInProcessServer(t)) + require.NoError(t, manager.AddClient(cfg)) + + monitor := mcp.NewClientHealthMonitor(manager, cfg.ID, 10*time.Millisecond, true, core.NewDefaultLogger(schemas.LogLevelError)) + monitor.Start() + defer monitor.Stop() + + // Wait for at least two complete cycles (both pre and post). + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if plugin.GetPreHookCallCount() >= 2 && plugin.GetPostHookCallCount() >= 2 { + break + } + time.Sleep(10 * time.Millisecond) + } + + require.GreaterOrEqual(t, plugin.GetPreHookCallCount(), 2, "PreHook fires even when short-circuiting") + // Short-circuited healthy → PostHook still runs with the synthetic response. + require.GreaterOrEqual(t, plugin.GetPostHookCallCount(), 2, "PostHook fires for short-circuit success path") +} + +func TestPingHook_PreHookShortCircuitError_DoesNotPanic(t *testing.T) { + t.Parallel() + + // Short-circuiting with error is treated by the health monitor as a normal + // ping failure. We verify the gate plumbing doesn't blow up and the plugin + // got invoked. + plugin := NewTestPingPlugin() + plugin.SetShortCircuitError(&schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "synthetic ping failure"}, + }) + + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin}) + cfg := inProcessClientConfig("ping_err", buildInProcessServer(t)) + require.NoError(t, manager.AddClient(cfg)) + + monitor := mcp.NewClientHealthMonitor(manager, cfg.ID, 10*time.Millisecond, true, core.NewDefaultLogger(schemas.LogLevelError)) + monitor.Start() + defer monitor.Stop() + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if plugin.GetPreHookCallCount() >= 2 { + break + } + time.Sleep(10 * time.Millisecond) + } + require.GreaterOrEqual(t, plugin.GetPreHookCallCount(), 2, "ping plugin fires regardless of short-circuit outcome") + // PostHook should NOT have captured anything because our plugin filters out + // non-healthy responses, and short-circuit-error skips the success path. + assert.Equal(t, 0, plugin.GetPostHookCallCount(), "no healthy ping response captured under error short-circuit") +} + +func TestPingHook_DoesNotFireWhenPingUnavailable(t *testing.T) { + t.Parallel() + + // When isPingAvailable=false, health monitor falls back to list_tools as the + // liveness probe. Ping hook should NEVER fire; list_tools hook fires instead. + pingPlugin := NewTestPingPlugin() + listPlugin := NewTestListToolsPlugin() + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{pingPlugin, listPlugin}) + + cfg := inProcessClientConfig("ping_unavailable", buildInProcessServer(t)) + require.NoError(t, manager.AddClient(cfg)) + + // Reset the list-tools plugin so we ignore the AddClient-time invocation. + listPlugin.Reset() + + monitor := mcp.NewClientHealthMonitor(manager, cfg.ID, 10*time.Millisecond, false, core.NewDefaultLogger(schemas.LogLevelError)) + monitor.Start() + defer monitor.Stop() + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if listPlugin.GetPreHookCallCount() >= 2 { + break + } + time.Sleep(10 * time.Millisecond) + } + + assert.Equal(t, 0, pingPlugin.GetPreHookCallCount(), "Ping hook MUST NOT fire when isPingAvailable=false") + assert.GreaterOrEqual(t, listPlugin.GetPreHookCallCount(), 2, "list_tools hook fires as liveness fallback") +} + +// ============================================================================= +// CROSS-CUTTING TESTS +// ============================================================================= + +func TestMCPGate_AllRequestTypesCarryClientName(t *testing.T) { + t.Parallel() + + // Verify ClientName is populated on requests handed to PreMCPHook for all + // three op types. Use the generic logging plugin so we see them all. + logPlugin := NewTestLoggingPlugin() + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{logPlugin}) + + require.NoError(t, manager.AddClient(inProcessClientConfig("client_name_test", buildInProcessServer(t)))) + + // Force a list_tools via reconnect to make sure we see at least one of each kind. + require.NoError(t, manager.ReconnectClient("client_name_test-id")) + + calls := logPlugin.GetPreHookCalls() + sawConnect := false + sawListTools := false + for _, c := range calls { + if c.ConnectRequest != nil { + // Typed Connect capture + assert.Equal(t, "client_name_test", c.ConnectRequest.ClientName, + "ClientName must be set on Connect sub-request") + sawConnect = true + continue + } + // Envelope capture (Ping / ListTools / ExecuteTool) + require.NotNil(t, c.Request) + assert.Equal(t, "client_name_test", c.Request.ClientName, + "ClientName must be set for %s requests", c.Request.RequestType) + if c.Request.RequestType == schemas.MCPRequestTypeListTools { + sawListTools = true + } + } + assert.True(t, sawConnect, "should have seen Connect requests via typed pipeline") + assert.True(t, sawListTools, "should have seen ListTools requests via envelope pipeline") +} + +func TestMCPGate_NoPluginsConfigured_OpStillRuns(t *testing.T) { + t.Parallel() + + // Even with no MCP plugins, the gate must transparently pass through. + manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{}) + + require.NoError(t, manager.AddClient(inProcessClientConfig("no_plugins", buildInProcessServer(t)))) + + clients := manager.GetClients() + var found bool + for _, c := range clients { + if c.Name == "no_plugins" { + found = true + assert.NotEmpty(t, c.ToolMap, "tools should have been discovered through the gate even without plugins") + break + } + } + assert.True(t, found, "client should be registered") +} diff --git a/core/internal/mcptests/context_propagation_test.go b/core/internal/mcptests/context_propagation_test.go index fb22dca4eff..c382f2d2500 100644 --- a/core/internal/mcptests/context_propagation_test.go +++ b/core/internal/mcptests/context_propagation_test.go @@ -259,9 +259,6 @@ func TestContext_ValueIsolation(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -398,9 +395,6 @@ func TestContext_TimeoutPropagation(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // May timeout or complete with partial results @@ -494,9 +488,6 @@ func TestContext_RequestIDGeneration(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) @@ -627,9 +618,6 @@ func TestContext_ConcurrentAccess(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) diff --git a/core/internal/mcptests/test_plugins.go b/core/internal/mcptests/test_plugins.go index 6178c4ca234..1b440f17be2 100644 --- a/core/internal/mcptests/test_plugins.go +++ b/core/internal/mcptests/test_plugins.go @@ -21,12 +21,17 @@ type TestLoggingPlugin struct { captureResponses bool } -// MCPLogEntry represents a logged MCP operation +// MCPLogEntry represents a logged MCP operation. For envelope-based ops +// (Ping/ListTools/ExecuteTool) Request/Response are populated. For typed Connect +// ops, ConnectRequest/ConnectResponse are populated instead — the two pipelines +// are separate so each entry carries exactly one shape. type MCPLogEntry struct { - Request *schemas.BifrostMCPRequest - Response *schemas.BifrostMCPResponse - Error *schemas.BifrostError - Timestamp int64 + Request *schemas.BifrostMCPRequest + Response *schemas.BifrostMCPResponse + ConnectRequest *schemas.BifrostMCPConnectRequest + ConnectResponse *schemas.BifrostMCPConnectResponse + Error *schemas.BifrostError + Timestamp int64 } // NewTestLoggingPlugin creates a new test logging plugin @@ -116,6 +121,35 @@ func (p *TestLoggingPlugin) Reset() { p.postHookCalls = make([]MCPLogEntry, 0) } +// PreMCPConnectionHook implements schemas.MCPConnectionPlugin so the logging plugin +// observes Connect events too. The typed sub-request lands in ConnectRequest on the +// log entry — the envelope-based Request field is left nil for Connect captures. +func (p *TestLoggingPlugin) PreMCPConnectionHook(ctx *schemas.BifrostContext, req *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectRequest, *schemas.MCPConnectionShortCircuit, error) { + if p.captureRequests { + p.mu.Lock() + p.preHookCalls = append(p.preHookCalls, MCPLogEntry{ + ConnectRequest: req, + Timestamp: time.Now().UnixNano(), + }) + p.mu.Unlock() + } + return req, nil, nil +} + +// PostMCPConnectionHook implements schemas.MCPConnectionPlugin for Connect responses. +func (p *TestLoggingPlugin) PostMCPConnectionHook(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPConnectResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError, error) { + if p.captureResponses { + p.mu.Lock() + p.postHookCalls = append(p.postHookCalls, MCPLogEntry{ + ConnectResponse: resp, + Error: bifrostErr, + Timestamp: time.Now().UnixNano(), + }) + p.mu.Unlock() + } + return resp, bifrostErr, nil +} + // ============================================================================= // TEST GOVERNANCE PLUGIN // ============================================================================= @@ -495,3 +529,436 @@ func (p *TestShortCircuitPlugin) PreMCPHook(ctx *schemas.BifrostContext, req *sc func (p *TestShortCircuitPlugin) PostMCPHook(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPResponse, *schemas.BifrostError, error) { return resp, bifrostErr, nil } + +// ============================================================================= +// TEST CONNECT PLUGIN — observes / mutates / short-circuits Connect requests +// ============================================================================= +// +// Only acts on Connect requests via the typed MCPConnectionPlugin interface. +// MCPPluginNoOpHooks provides no-op generic PreMCPHook/PostMCPHook so this +// plugin satisfies MCPPlugin (required by the BifrostConfig.MCPPlugins slice). +type TestConnectPlugin struct { + schemas.MCPPluginNoOpHooks + + mu sync.RWMutex + preHookCalls []MCPLogEntry + postHookCalls []MCPLogEntry + + // Mutation knobs (applied in PreHook if set). + mutateHeaders map[string]string + mutateConnString *string + mutateStdioCommand *string + mutateStdioArgs []string + mutateStdioArgsIsSet bool + + // Short-circuit knobs (PreHook). + shortCircuitResponse *schemas.BifrostMCPConnectResponse + shortCircuitError *schemas.BifrostError +} + +func NewTestConnectPlugin() *TestConnectPlugin { + return &TestConnectPlugin{ + preHookCalls: make([]MCPLogEntry, 0), + postHookCalls: make([]MCPLogEntry, 0), + } +} + +func (p *TestConnectPlugin) GetName() string { return "TestConnectPlugin" } +func (p *TestConnectPlugin) Cleanup() error { return nil } + +// SetMutateHeaders configures the plugin to overwrite the Headers field in PreHook. +func (p *TestConnectPlugin) SetMutateHeaders(headers map[string]string) { + p.mu.Lock() + defer p.mu.Unlock() + p.mutateHeaders = headers +} + +// SetMutateConnectionString configures the plugin to overwrite ConnectionString in PreHook. +func (p *TestConnectPlugin) SetMutateConnectionString(url *string) { + p.mu.Lock() + defer p.mu.Unlock() + p.mutateConnString = url +} + +// SetMutateStdioCommand configures the plugin to overwrite StdioCommand in PreHook. +func (p *TestConnectPlugin) SetMutateStdioCommand(cmd *string) { + p.mu.Lock() + defer p.mu.Unlock() + p.mutateStdioCommand = cmd +} + +// SetMutateStdioArgs configures the plugin to overwrite StdioArgs in PreHook. +// Pass nil to leave it unchanged; pass [] to explicitly clear it. +func (p *TestConnectPlugin) SetMutateStdioArgs(args []string) { + p.mu.Lock() + defer p.mu.Unlock() + p.mutateStdioArgs = args + p.mutateStdioArgsIsSet = true +} + +// SetShortCircuitResponse configures the plugin to short-circuit Connect with a +// synthetic success response carrying the provided sub-response payload. +func (p *TestConnectPlugin) SetShortCircuitResponse(resp *schemas.BifrostMCPConnectResponse) { + p.mu.Lock() + defer p.mu.Unlock() + p.shortCircuitResponse = resp +} + +// SetShortCircuitError configures the plugin to short-circuit Connect with the given error. +func (p *TestConnectPlugin) SetShortCircuitError(err *schemas.BifrostError) { + p.mu.Lock() + defer p.mu.Unlock() + p.shortCircuitError = err +} + +// PreMCPConnectionHook implements schemas.MCPConnectionPlugin (typed Connect hook). +// No RequestType filtering needed — the pipeline only invokes this method for +// Connect requests, and it gets the typed sub-request directly. +func (p *TestConnectPlugin) PreMCPConnectionHook(ctx *schemas.BifrostContext, req *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectRequest, *schemas.MCPConnectionShortCircuit, error) { + p.mu.Lock() + defer p.mu.Unlock() + + p.preHookCalls = append(p.preHookCalls, MCPLogEntry{ + ConnectRequest: req, + Timestamp: time.Now().UnixNano(), + }) + + // Short-circuit before mutation (mutation only matters if the op runs). + if p.shortCircuitError != nil { + return req, &schemas.MCPConnectionShortCircuit{Error: p.shortCircuitError}, nil + } + if p.shortCircuitResponse != nil { + return req, &schemas.MCPConnectionShortCircuit{Response: p.shortCircuitResponse}, nil + } + + // Apply mutations directly on the typed sub-request — no nil-check on a wrapper needed. + if p.mutateHeaders != nil { + req.Headers = p.mutateHeaders + } + if p.mutateConnString != nil { + req.ConnectionString = p.mutateConnString + } + if p.mutateStdioCommand != nil { + req.StdioCommand = p.mutateStdioCommand + } + if p.mutateStdioArgsIsSet { + req.StdioArgs = p.mutateStdioArgs + } + return req, nil, nil +} + +// PostMCPConnectionHook implements schemas.MCPConnectionPlugin. +// Captures only successful Connect outcomes (resp non-nil). Short-circuit-error +// paths skip capture — matching the "observe outcomes" intent of test logging. +func (p *TestConnectPlugin) PostMCPConnectionHook(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPConnectResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError, error) { + if resp == nil { + return resp, bifrostErr, nil + } + p.mu.Lock() + p.postHookCalls = append(p.postHookCalls, MCPLogEntry{ + ConnectResponse: resp, + Error: bifrostErr, + Timestamp: time.Now().UnixNano(), + }) + p.mu.Unlock() + return resp, bifrostErr, nil +} + +func (p *TestConnectPlugin) GetPreHookCalls() []MCPLogEntry { + p.mu.RLock() + defer p.mu.RUnlock() + out := make([]MCPLogEntry, len(p.preHookCalls)) + copy(out, p.preHookCalls) + return out +} + +func (p *TestConnectPlugin) GetPostHookCalls() []MCPLogEntry { + p.mu.RLock() + defer p.mu.RUnlock() + out := make([]MCPLogEntry, len(p.postHookCalls)) + copy(out, p.postHookCalls) + return out +} + +func (p *TestConnectPlugin) Reset() { + p.mu.Lock() + defer p.mu.Unlock() + p.preHookCalls = p.preHookCalls[:0] + p.postHookCalls = p.postHookCalls[:0] +} + +// ============================================================================= +// TEST PING PLUGIN — observes / short-circuits Ping requests +// ============================================================================= +// +// Only acts on requests with RequestType == MCPRequestTypePing. +type TestPingPlugin struct { + mu sync.RWMutex + preHookCalls []MCPLogEntry + postHookCalls []MCPLogEntry + + shortCircuitHealthy bool // if true, PreHook returns a synthetic healthy response + shortCircuitError *schemas.BifrostError // if non-nil, PreHook returns this error +} + +func NewTestPingPlugin() *TestPingPlugin { + return &TestPingPlugin{ + preHookCalls: make([]MCPLogEntry, 0), + postHookCalls: make([]MCPLogEntry, 0), + } +} + +func (p *TestPingPlugin) GetName() string { return "TestPingPlugin" } +func (p *TestPingPlugin) Cleanup() error { return nil } + +// SetShortCircuitHealthy makes PreHook return a synthetic healthy ping response. +func (p *TestPingPlugin) SetShortCircuitHealthy(healthy bool) { + p.mu.Lock() + defer p.mu.Unlock() + p.shortCircuitHealthy = healthy +} + +// SetShortCircuitError makes PreHook return the given error (counts as ping failure). +func (p *TestPingPlugin) SetShortCircuitError(err *schemas.BifrostError) { + p.mu.Lock() + defer p.mu.Unlock() + p.shortCircuitError = err +} + +func (p *TestPingPlugin) PreMCPHook(ctx *schemas.BifrostContext, req *schemas.BifrostMCPRequest) (*schemas.BifrostMCPRequest, *schemas.MCPPluginShortCircuit, error) { + if req == nil || req.RequestType != schemas.MCPRequestTypePing { + return req, nil, nil + } + + p.mu.Lock() + defer p.mu.Unlock() + + p.preHookCalls = append(p.preHookCalls, MCPLogEntry{ + Request: req, + Timestamp: time.Now().UnixNano(), + }) + + if p.shortCircuitError != nil { + return req, &schemas.MCPPluginShortCircuit{Error: p.shortCircuitError}, nil + } + if p.shortCircuitHealthy { + return req, &schemas.MCPPluginShortCircuit{ + Response: &schemas.BifrostMCPResponse{ + BifrostMCPPingResponse: &schemas.BifrostMCPPingResponse{}, + }, + }, nil + } + return req, nil, nil +} + +func (p *TestPingPlugin) PostMCPHook(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPResponse, *schemas.BifrostError, error) { + // Distinguish ping responses from other ops. Successful ping carries a non-nil + // BifrostMCPPingResponse; failed ping has nil response + non-nil error — in that + // case we can't tell from the response alone, but the err path is reached for + // any failed op, so for now only capture successful pings (matches the typical + // observability use case). + if resp == nil || resp.BifrostMCPPingResponse == nil { + return resp, bifrostErr, nil + } + p.mu.Lock() + p.postHookCalls = append(p.postHookCalls, MCPLogEntry{ + Response: resp, + Error: bifrostErr, + Timestamp: time.Now().UnixNano(), + }) + p.mu.Unlock() + return resp, bifrostErr, nil +} + +func (p *TestPingPlugin) GetPreHookCallCount() int { + p.mu.RLock() + defer p.mu.RUnlock() + return len(p.preHookCalls) +} + +func (p *TestPingPlugin) GetPostHookCallCount() int { + p.mu.RLock() + defer p.mu.RUnlock() + return len(p.postHookCalls) +} + +func (p *TestPingPlugin) GetPreHookCalls() []MCPLogEntry { + p.mu.RLock() + defer p.mu.RUnlock() + out := make([]MCPLogEntry, len(p.preHookCalls)) + copy(out, p.preHookCalls) + return out +} + +func (p *TestPingPlugin) GetPostHookCalls() []MCPLogEntry { + p.mu.RLock() + defer p.mu.RUnlock() + out := make([]MCPLogEntry, len(p.postHookCalls)) + copy(out, p.postHookCalls) + return out +} + +func (p *TestPingPlugin) Reset() { + p.mu.Lock() + defer p.mu.Unlock() + p.preHookCalls = p.preHookCalls[:0] + p.postHookCalls = p.postHookCalls[:0] +} + +// ============================================================================= +// TEST LISTTOOLS PLUGIN — observes / mutates / short-circuits ListTools requests +// ============================================================================= +// +// Only acts on requests with RequestType == MCPRequestTypeListTools. +type TestListToolsPlugin struct { + mu sync.RWMutex + preHookCalls []MCPLogEntry + postHookCalls []MCPLogEntry + + // PreHook short-circuit knobs. + shortCircuitResponse *schemas.BifrostMCPListToolsResponse + shortCircuitError *schemas.BifrostError + + // PostHook mutation knob: optional filter applied to the Tools map. If non-nil, + // only keys returned true are kept; ToolNameMapping is filtered to match. + postHookFilter func(toolName string) bool +} + +func NewTestListToolsPlugin() *TestListToolsPlugin { + return &TestListToolsPlugin{ + preHookCalls: make([]MCPLogEntry, 0), + postHookCalls: make([]MCPLogEntry, 0), + } +} + +func (p *TestListToolsPlugin) GetName() string { return "TestListToolsPlugin" } +func (p *TestListToolsPlugin) Cleanup() error { return nil } + +// SetShortCircuitResponse makes PreHook return a synthetic list_tools response. +func (p *TestListToolsPlugin) SetShortCircuitResponse(resp *schemas.BifrostMCPListToolsResponse) { + p.mu.Lock() + defer p.mu.Unlock() + p.shortCircuitResponse = resp +} + +// SetShortCircuitError makes PreHook return the given error. +func (p *TestListToolsPlugin) SetShortCircuitError(err *schemas.BifrostError) { + p.mu.Lock() + defer p.mu.Unlock() + p.shortCircuitError = err +} + +// SetPostHookFilter configures a PostHook tool-name predicate. Tools whose prefixed +// name passes (returns true) are kept; others are removed from both Tools and +// ToolNameMapping (matching the sanitized->original lookup). +func (p *TestListToolsPlugin) SetPostHookFilter(filter func(toolName string) bool) { + p.mu.Lock() + defer p.mu.Unlock() + p.postHookFilter = filter +} + +func (p *TestListToolsPlugin) PreMCPHook(ctx *schemas.BifrostContext, req *schemas.BifrostMCPRequest) (*schemas.BifrostMCPRequest, *schemas.MCPPluginShortCircuit, error) { + if req == nil || req.RequestType != schemas.MCPRequestTypeListTools { + return req, nil, nil + } + + p.mu.Lock() + defer p.mu.Unlock() + + p.preHookCalls = append(p.preHookCalls, MCPLogEntry{ + Request: req, + Timestamp: time.Now().UnixNano(), + }) + + if p.shortCircuitError != nil { + return req, &schemas.MCPPluginShortCircuit{Error: p.shortCircuitError}, nil + } + if p.shortCircuitResponse != nil { + return req, &schemas.MCPPluginShortCircuit{ + Response: &schemas.BifrostMCPResponse{ + BifrostMCPListToolsResponse: p.shortCircuitResponse, + }, + }, nil + } + return req, nil, nil +} + +func (p *TestListToolsPlugin) PostMCPHook(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPResponse, *schemas.BifrostError, error) { + if resp == nil || resp.BifrostMCPListToolsResponse == nil { + return resp, bifrostErr, nil + } + + p.mu.Lock() + defer p.mu.Unlock() + + // Apply filter if configured. Tools are keyed by prefixed name; ToolNameMapping + // is keyed by sanitized name — we filter by prefixed name and drop the matching + // sanitized entry too (sanitized = stripped client prefix, with '-'→'_'). + if p.postHookFilter != nil { + filteredTools := make(map[string]schemas.ChatTool, len(resp.Tools)) + keep := make(map[string]bool, len(resp.Tools)) + for name, tool := range resp.Tools { + if p.postHookFilter(name) { + filteredTools[name] = tool + keep[name] = true + } + } + // Filter ToolNameMapping: keep mapping entries whose original value still + // corresponds to a tool that survived. + filteredMapping := make(map[string]string, len(resp.ToolNameMapping)) + for sanitized, original := range resp.ToolNameMapping { + // Find the prefixed key that would have been used for `original`. + for prefixed := range keep { + // prefixed == "-"; check suffix match. + if len(prefixed) > len(original) && prefixed[len(prefixed)-len(original):] == original { + filteredMapping[sanitized] = original + break + } + } + } + resp.Tools = filteredTools + resp.ToolNameMapping = filteredMapping + } + + p.postHookCalls = append(p.postHookCalls, MCPLogEntry{ + Response: resp, + Error: bifrostErr, + Timestamp: time.Now().UnixNano(), + }) + return resp, bifrostErr, nil +} + +func (p *TestListToolsPlugin) GetPreHookCallCount() int { + p.mu.RLock() + defer p.mu.RUnlock() + return len(p.preHookCalls) +} + +func (p *TestListToolsPlugin) GetPostHookCallCount() int { + p.mu.RLock() + defer p.mu.RUnlock() + return len(p.postHookCalls) +} + +func (p *TestListToolsPlugin) GetPreHookCalls() []MCPLogEntry { + p.mu.RLock() + defer p.mu.RUnlock() + out := make([]MCPLogEntry, len(p.preHookCalls)) + copy(out, p.preHookCalls) + return out +} + +func (p *TestListToolsPlugin) GetPostHookCalls() []MCPLogEntry { + p.mu.RLock() + defer p.mu.RUnlock() + out := make([]MCPLogEntry, len(p.postHookCalls)) + copy(out, p.postHookCalls) + return out +} + +func (p *TestListToolsPlugin) Reset() { + p.mu.Lock() + defer p.mu.Unlock() + p.preHookCalls = p.preHookCalls[:0] + p.postHookCalls = p.postHookCalls[:0] +} diff --git a/core/internal/mcptests/tool_call_id_test.go b/core/internal/mcptests/tool_call_id_test.go index beafa63d481..e2be70cc0f5 100644 --- a/core/internal/mcptests/tool_call_id_test.go +++ b/core/internal/mcptests/tool_call_id_test.go @@ -154,9 +154,6 @@ func TestToolCallID_DuplicateIDsInParallelExecution(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) // Should complete (even if results might be ambiguous) @@ -327,9 +324,6 @@ func TestToolCallID_UniqueIDsInBatch(t *testing.T) { originalReq, initialResponse, mockLLM.MakeChatRequest, - func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { - return manager.ExecuteToolCall(ctx, request) - }, ) require.Nil(t, bifrostErr) diff --git a/core/mcp/agent.go b/core/mcp/agent.go index 96d16ec24ee..1aa7615097e 100644 --- a/core/mcp/agent.go +++ b/core/mcp/agent.go @@ -317,11 +317,9 @@ func (a *AgentModeExecutor) executeAgent( channelToolResults <- createToolResultMessage(toolCall, "", toolErr) } else if mcpResponse != nil && mcpResponse.ChatMessage != nil { channelToolResults <- mcpResponse.ChatMessage - } else if mcpResponse != nil && mcpResponse.ChatMessage == nil { - // Send empty result when mcpResponse is non-nil but ChatMessage is nil - channelToolResults <- createToolResultMessage(toolCall, "", nil) } else { - // Fallback: send empty result when both mcpResponse and toolErr are nil + // Fallback: empty result when mcpResponse is missing the chat message + // (either nil mcpResponse, missing execute-tool payload, or nil ChatMessage). channelToolResults <- createToolResultMessage(toolCall, "", nil) } }(toolCall) diff --git a/core/mcp/agent_test.go b/core/mcp/agent_test.go index fb44451a6d6..e1b4b11bd51 100644 --- a/core/mcp/agent_test.go +++ b/core/mcp/agent_test.go @@ -77,6 +77,9 @@ func (m *MockClientManager) GetToolPerClient(ctx context.Context) map[string][]s return make(map[string][]schemas.ChatTool) } +func (m *MockClientManager) GetPluginPipeline() PluginPipeline { return nil } +func (m *MockClientManager) ReleasePluginPipeline(pipeline PluginPipeline) {} + func TestHasToolCallsForChatResponse(t *testing.T) { // Test nil response if hasToolCallsForChatResponse(nil) { @@ -556,6 +559,9 @@ func (m *MockAutoClientManager) GetToolPerClient(ctx context.Context) map[string return make(map[string][]schemas.ChatTool) } +func (m *MockAutoClientManager) GetPluginPipeline() PluginPipeline { return nil } +func (m *MockAutoClientManager) ReleasePluginPipeline(pipeline PluginPipeline) {} + // TestParallelToolCallsHaveUniqueMCPLogIDs verifies that parallel tool calls within a // single LLM response each receive a unique BifrostContextKeyMCPLogID in their context. // diff --git a/core/mcp/clientmanager.go b/core/mcp/clientmanager.go index c1cfb15af5a..1a83693a0b0 100644 --- a/core/mcp/clientmanager.go +++ b/core/mcp/clientmanager.go @@ -7,6 +7,7 @@ import ( "os" "slices" "strings" + "time" "github.com/mark3labs/mcp-go/client" "github.com/mark3labs/mcp-go/client/transport" @@ -208,43 +209,115 @@ func (m *MCPManager) VerifyPerUserOAuthConnection(ctx context.Context, config *s return nil, nil, fmt.Errorf("connection URL is required for per-user OAuth verification") } - // Create HTTP transport with the admin's temporary Bearer token - headers := map[string]string{ - "Authorization": "Bearer " + accessToken, + // Build prepared inputs for the typed connect plugin gate. PreHooks may mutate + // Headers / ConnectionString — the mutated values are passed to the transport below. + // Copy non-Authorization headers from config.Headers so verification sees the same + // tenant/custom headers as the normal connect path. Authorization is re-injected + // after PreHooks run (the OAuth bearer comes from the access token, not config). + url := config.ConnectionString.GetValue() + preparedHeaders := make(map[string]string, len(config.Headers)) + for k, v := range config.Headers { + if strings.EqualFold(k, "Authorization") { + continue + } + preparedHeaders[k] = v.GetValue() } - httpTransport, err := transport.NewStreamableHTTP(config.ConnectionString.GetValue(), transport.WithHTTPHeaders(headers)) - if err != nil { - return nil, nil, fmt.Errorf("failed to create HTTP transport for verification: %w", err) + connectReq := &schemas.BifrostMCPConnectRequest{ + ClientName: config.Name, + ConnectionType: schemas.MCPConnectionTypeHTTP, + AuthType: config.AuthType, + ConnectionString: &url, + Headers: preparedHeaders, } - // Create temporary MCP client - tempClient := client.NewClient(httpTransport) - ctx, cancel := context.WithTimeout(ctx, MCPClientConnectionEstablishTimeout) + verifyCtx, cancel := context.WithTimeout(ctx, MCPClientConnectionEstablishTimeout) defer cancel() + gateCtx := schemas.NewBifrostContext(verifyCtx, schemas.NoDeadline) - // Start transport - if err := tempClient.Start(ctx); err != nil { - return nil, nil, fmt.Errorf("failed to start MCP connection for verification: %w", err) - } - defer tempClient.Close() + var tempClient *client.Client + defer func() { + if tempClient != nil { + tempClient.Close() + } + }() + start := time.Now() + + _, gateErr := m.runConnectWithPluginPipeline(gateCtx, connectReq, func(preReq *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectResponse, error) { + // Use mutated URL/headers + finalURL := url + if preReq.ConnectionString != nil { + finalURL = *preReq.ConnectionString + } - // Initialize MCP handshake - initRequest := mcp.InitializeRequest{ - Params: mcp.InitializeParams{ - ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, - Capabilities: mcp.ClientCapabilities{}, - ClientInfo: mcp.Implementation{ - Name: fmt.Sprintf("Bifrost-%s-verify", config.Name), - Version: "1.0.0", + // Copy mutated headers and add Authorization AFTER all PreHooks ran. Copying + // (rather than mutating preReq.Headers in place) avoids leaking the bearer token + // back into the request that PreHook plugins may still reference. + finalHeaders := make(map[string]string, len(preReq.Headers)+1) + maps.Copy(finalHeaders, preReq.Headers) + finalHeaders["Authorization"] = fmt.Sprintf("Bearer %s", accessToken) + + httpTransport, hErr := transport.NewStreamableHTTP(finalURL, transport.WithHTTPHeaders(finalHeaders)) + if hErr != nil { + return nil, fmt.Errorf("failed to create HTTP transport for verification: %w", hErr) + } + tempClient = client.NewClient(httpTransport) + if startErr := tempClient.Start(verifyCtx); startErr != nil { + return nil, fmt.Errorf("failed to start MCP connection for verification: %w", startErr) + } + + initRequest := mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + Capabilities: mcp.ClientCapabilities{}, + ClientInfo: mcp.Implementation{ + Name: fmt.Sprintf("Bifrost-%s-verify", config.Name), + Version: "1.0.0", + }, }, - }, + } + initResult, initErr := tempClient.Initialize(verifyCtx, initRequest) + if initErr != nil { + return nil, fmt.Errorf("failed to initialize MCP connection for verification: %w", initErr) + } + + resp := &schemas.BifrostMCPConnectResponse{ + ConnectionInfo: &schemas.MCPClientConnectionInfo{ + Type: schemas.MCPConnectionTypeHTTP, + ConnectionURL: &finalURL, + }, + ExtraFields: schemas.BifrostMCPResponseExtraFields{ + Latency: time.Since(start).Milliseconds(), + }, + } + if initResult != nil { + resp.ProtocolVersion = initResult.ProtocolVersion + resp.ServerInfo = &schemas.MCPServerInfo{ + Name: initResult.ServerInfo.Name, + Version: initResult.ServerInfo.Version, + } + resp.ServerCapabilities = &schemas.MCPServerCapabilities{ + Tools: initResult.Capabilities.Tools != nil, + Resources: initResult.Capabilities.Resources != nil, + Prompts: initResult.Capabilities.Prompts != nil, + Logging: initResult.Capabilities.Logging != nil, + } + } + return resp, nil + }) + + if gateErr != nil { + return nil, nil, fmt.Errorf("failed to verify MCP connection: %s", gateErr.GetErrorString()) } - if _, err := tempClient.Initialize(ctx, initRequest); err != nil { - return nil, nil, fmt.Errorf("failed to initialize MCP connection for verification: %w", err) + if tempClient == nil { + // Plugin short-circuited connect with a synthetic success response. We have no live + // socket to query for tools — surface this as an error since tool discovery is the + // whole point of OAuth verification. + return nil, nil, fmt.Errorf("OAuth verification was short-circuited by plugin; cannot discover tools without a live connection") } - // Discover tools - tools, toolNameMapping, err := retrieveExternalTools(ctx, tempClient, config.Name, m.logger) + // Discover tools through the list_tools plugin gate. PostHook may filter or augment + // the discovered set. + tools, toolNameMapping, err := m.runListToolsWithHooks(verifyCtx, tempClient, config.Name) if err != nil { return nil, nil, fmt.Errorf("failed to discover tools during verification: %w", err) } @@ -834,7 +907,6 @@ func (m *MCPManager) connectToMCPClient(config *schemas.MCPClientConfig) error { // Heavy operations performed outside lock var externalClient *client.Client var connectionInfo *schemas.MCPClientConnectionInfo - var err error // Initialize the external client with timeout // For SSE and STDIO connections, we need a long-lived context for the connection @@ -860,134 +932,210 @@ func (m *MCPManager) connectToMCPClient(config *schemas.MCPClientConfig) error { defer cancel() } - // Start the transport first (required for STDIO and SSE clients) with retry logic - // Each retry attempt uses a fresh client instance to avoid resource leaks - m.logger.Debug("%s [%s] Starting transport...", MCPLogPrefix, config.Name) - transportRetryConfig := DefaultRetryConfig - err = ExecuteWithRetry( - m.ctx, - func() error { - // Close previous client if this is a retry attempt - if externalClient != nil { - if closeErr := externalClient.Close(); closeErr != nil { - m.logger.Warn("%s Failed to close external client during retry: %v", MCPLogPrefix, closeErr) - } - } - // Create a fresh client for this attempt - var createErr error - switch config.ConnectionType { - case schemas.MCPConnectionTypeHTTP: - externalClient, connectionInfo, createErr = m.createHTTPConnection(m.ctx, config) - case schemas.MCPConnectionTypeSTDIO: - externalClient, connectionInfo, createErr = m.createSTDIOConnection(m.ctx, config) - case schemas.MCPConnectionTypeSSE: - externalClient, connectionInfo, createErr = m.createSSEConnection(m.ctx, config) - case schemas.MCPConnectionTypeInProcess: - externalClient, connectionInfo, createErr = m.createInProcessConnection(m.ctx, config) - default: - return fmt.Errorf("unknown connection type: %s", config.ConnectionType) + // Build the plugin gate request with prepared inputs. PreHooks may mutate + // ConnectionString, Headers, StdioCommand, StdioArgs — those mutations flow into + // the createConnection calls below via the `overrides` parameter. + // + // SECURITY: Authorization is stripped from the headers exposed to PreHooks and + // re-injected after all PreHooks have run. Plugins never see the bearer token. + connectReq := &schemas.BifrostMCPConnectRequest{ + ClientName: config.Name, + ConnectionType: config.ConnectionType, + AuthType: config.AuthType, + } + if config.ConnectionString != nil { + u := config.ConnectionString.GetValue() + connectReq.ConnectionString = &u + } + var authHeader string // captured for re-injection after PreHooks + if config.ConnectionType == schemas.MCPConnectionTypeHTTP || config.ConnectionType == schemas.MCPConnectionTypeSSE { + if h, hErr := config.HttpHeaders(m.ctx, m.oauth2Provider); hErr == nil { + if auth, ok := h["Authorization"]; ok { + authHeader = auth + delete(h, "Authorization") } - if createErr != nil { - return createErr - } - // Create per-attempt timeout context for Start operation - // Each attempt has a deadline to prevent indefinite hangs - var perAttemptCtx context.Context - if config.ConnectionType == schemas.MCPConnectionTypeSSE || config.ConnectionType == schemas.MCPConnectionTypeSTDIO { - // For STDIO/SSE: use longLivedCtx directly without additional timeout - // The subprocess needs the context to stay valid for the entire connection lifetime - // Do NOT defer cancel - the context manages the subprocess lifetime - perAttemptCtx = longLivedCtx - m.logger.Debug("%s [%s] Starting transport...", MCPLogPrefix, config.Name) - } else { - // HTTP already has timeout - perAttemptCtx = ctx - } - // Start the fresh client with the per-attempt timeout - return externalClient.Start(perAttemptCtx) - }, - transportRetryConfig, - m.logger, - ) - if err != nil { - if config.ConnectionType == schemas.MCPConnectionTypeSSE || config.ConnectionType == schemas.MCPConnectionTypeSTDIO { - cancel() // Cancel long-lived context on error + connectReq.Headers = h } - // Close external client connection to prevent transport/goroutine leaks - if externalClient != nil { - if closeErr := externalClient.Close(); closeErr != nil { - m.logger.Warn("%s Failed to close external client during cleanup: %v", MCPLogPrefix, closeErr) - } - } - return fmt.Errorf("failed to start MCP client transport %s after %d retries: %v", config.Name, transportRetryConfig.MaxRetries, err) } - m.logger.Debug("%s [%s] Transport started successfully", MCPLogPrefix, config.Name) + if config.StdioConfig != nil { + cmd := config.StdioConfig.Command + connectReq.StdioCommand = &cmd + connectReq.StdioArgs = append([]string(nil), config.StdioConfig.Args...) + } + + // Fresh BifrostContext for the gate so it doesn't inherit any SkipPluginPipeline flag + // from the caller's request context. Connect runs as infrastructure, not as part of an + // in-flight LLM request. + gateCtx := schemas.NewBifrostContext(m.ctx, schemas.NoDeadline) + + // To capture InitializeResult for the response, the op closure populates these. + var initResult *mcp.InitializeResult + start := time.Now() + + _, gateErr := m.runConnectWithPluginPipeline(gateCtx, connectReq, func(preReq *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectResponse, error) { + // Re-inject Authorization after PreHooks. Use a shallow-cloned overrides + // struct so the merged headers don't leak back into the request object that + // plugins captured in PreHook (which they may still reference in PostHook). + mutForWire := preReq + if authHeader != "" { + merged := make(map[string]string, len(preReq.Headers)+1) + maps.Copy(merged, preReq.Headers) + merged["Authorization"] = authHeader + clone := *preReq + clone.Headers = merged + mutForWire = &clone + } - // Create proper initialize request for external client - extInitRequest := mcp.InitializeRequest{ - Params: mcp.InitializeParams{ - ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, - Capabilities: mcp.ClientCapabilities{}, - ClientInfo: mcp.Implementation{ - Name: fmt.Sprintf("Bifrost-%s", config.Name), - Version: "1.0.0", + // Start the transport (with internal retries). Each retry uses a fresh client. + m.logger.Debug("%s [%s] Starting transport...", MCPLogPrefix, config.Name) + transportRetryConfig := DefaultRetryConfig + if startErr := ExecuteWithRetry( + m.ctx, + func() error { + // Close previous client if this is a retry attempt + if externalClient != nil { + if closeErr := externalClient.Close(); closeErr != nil { + m.logger.Warn("%s Failed to close external client during retry: %v", MCPLogPrefix, closeErr) + } + } + // Create a fresh client for this attempt + var createErr error + switch config.ConnectionType { + case schemas.MCPConnectionTypeHTTP: + externalClient, connectionInfo, createErr = m.createHTTPConnection(m.ctx, config, mutForWire) + case schemas.MCPConnectionTypeSTDIO: + externalClient, connectionInfo, createErr = m.createSTDIOConnection(m.ctx, config, mutForWire) + case schemas.MCPConnectionTypeSSE: + externalClient, connectionInfo, createErr = m.createSSEConnection(m.ctx, config, mutForWire) + case schemas.MCPConnectionTypeInProcess: + externalClient, connectionInfo, createErr = m.createInProcessConnection(m.ctx, config) + default: + return fmt.Errorf("unknown connection type: %s", config.ConnectionType) + } + if createErr != nil { + return createErr + } + // Create per-attempt timeout context for Start operation + // Each attempt has a deadline to prevent indefinite hangs + var perAttemptCtx context.Context + if config.ConnectionType == schemas.MCPConnectionTypeSSE || config.ConnectionType == schemas.MCPConnectionTypeSTDIO { + // For STDIO/SSE: use longLivedCtx directly without additional timeout + // The subprocess needs the context to stay valid for the entire connection lifetime + // Do NOT defer cancel - the context manages the subprocess lifetime. + perAttemptCtx = longLivedCtx + m.logger.Debug("%s [%s] Starting transport...", MCPLogPrefix, config.Name) + } else { + // HTTP already has timeout + perAttemptCtx = ctx + } + return externalClient.Start(perAttemptCtx) }, - }, - } + transportRetryConfig, + m.logger, + ); startErr != nil { + return nil, fmt.Errorf("failed to start MCP client transport after %d retries: %v", transportRetryConfig.MaxRetries, startErr) + } + m.logger.Debug("%s [%s] Transport started successfully", MCPLogPrefix, config.Name) + + // Initialize with retry. Capture InitializeResult so the gate response can expose + // ServerInfo / ProtocolVersion / Capabilities. + extInitRequest := mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + Capabilities: mcp.ClientCapabilities{}, + ClientInfo: mcp.Implementation{ + Name: fmt.Sprintf("Bifrost-%s", config.Name), + Version: "1.0.0", + }, + }, + } + initRetryConfig := DefaultRetryConfig + if initErr := ExecuteWithRetry( + m.ctx, + func() error { + var initCtx context.Context + if config.ConnectionType == schemas.MCPConnectionTypeSSE || config.ConnectionType == schemas.MCPConnectionTypeSTDIO { + var initCancel context.CancelFunc + initCtx, initCancel = context.WithTimeout(longLivedCtx, MCPClientConnectionEstablishTimeout) + defer initCancel() + m.logger.Debug("%s [%s] Initializing client with %v timeout...", MCPLogPrefix, config.Name, MCPClientConnectionEstablishTimeout) + } else { + initCtx = ctx + } + var initErr error + initResult, initErr = externalClient.Initialize(initCtx, extInitRequest) + return initErr + }, + initRetryConfig, + m.logger, + ); initErr != nil { + return nil, fmt.Errorf("failed to initialize MCP client after %d retries: %v", initRetryConfig.MaxRetries, initErr) + } + m.logger.Debug("%s [%s] Client initialized successfully", MCPLogPrefix, config.Name) - // Initialize client with retry logic - initRetryConfig := DefaultRetryConfig - err = ExecuteWithRetry( - m.ctx, - func() error { - // For STDIO/SSE: Use a timeout context for initialization to prevent indefinite hangs - // The subprocess will continue running with the long-lived context - var initCtx context.Context - var initCancel context.CancelFunc - - if config.ConnectionType == schemas.MCPConnectionTypeSSE || config.ConnectionType == schemas.MCPConnectionTypeSTDIO { - // Create timeout context for initialization phase only - initCtx, initCancel = context.WithTimeout(longLivedCtx, MCPClientConnectionEstablishTimeout) - defer initCancel() - m.logger.Debug("%s [%s] Initializing client with %v timeout...", MCPLogPrefix, config.Name, MCPClientConnectionEstablishTimeout) - } else { - // HTTP already has timeout - initCtx = ctx + // Build the gate response from captured initialize result. + resp := &schemas.BifrostMCPConnectResponse{ + ConnectionInfo: connectionInfo, + ExtraFields: schemas.BifrostMCPResponseExtraFields{ + Latency: time.Since(start).Milliseconds(), + }, + } + if initResult != nil { + resp.ProtocolVersion = initResult.ProtocolVersion + resp.ServerInfo = &schemas.MCPServerInfo{ + Name: initResult.ServerInfo.Name, + Version: initResult.ServerInfo.Version, } - _, initErr := externalClient.Initialize(initCtx, extInitRequest) - return initErr - }, - initRetryConfig, - m.logger, - ) - if err != nil { + resp.ServerCapabilities = &schemas.MCPServerCapabilities{ + Tools: initResult.Capabilities.Tools != nil, + Resources: initResult.Capabilities.Resources != nil, + Prompts: initResult.Capabilities.Prompts != nil, + Logging: initResult.Capabilities.Logging != nil, + } + } + return resp, nil + }) + + if gateErr != nil { if config.ConnectionType == schemas.MCPConnectionTypeSSE || config.ConnectionType == schemas.MCPConnectionTypeSTDIO { - cancel() // Cancel long-lived context on error + cancel() } - // Close external client connection to prevent transport/goroutine leaks if externalClient != nil { if closeErr := externalClient.Close(); closeErr != nil { m.logger.Warn("%s Failed to close external client during cleanup: %v", MCPLogPrefix, closeErr) } } - return fmt.Errorf("failed to initialize MCP client %s after %d retries: %v", config.Name, initRetryConfig.MaxRetries, err) - } - m.logger.Debug("%s [%s] Client initialized successfully", MCPLogPrefix, config.Name) - - // Retrieve tools from the external server (this also requires network I/O) - // Use a bounded timeout context to prevent indefinite hangs during tool retrieval. - // For STDIO/SSE, ctx is longLivedCtx (no timeout), so we create a separate one here. - m.logger.Debug("%s [%s] Retrieving tools...", MCPLogPrefix, config.Name) - toolRetrievalCtx, toolRetrievalCancel := context.WithTimeout(m.ctx, MCPClientConnectionEstablishTimeout) - defer toolRetrievalCancel() - tools, toolNameMapping, err := retrieveExternalTools(toolRetrievalCtx, externalClient, config.Name, m.logger) - if err != nil { - m.logger.Warn("%s Failed to retrieve tools from %s: %v", MCPLogPrefix, config.Name, err) - // Continue with connection even if tool retrieval fails - tools = make(map[string]schemas.ChatTool) - toolNameMapping = make(map[string]string) + return fmt.Errorf("failed to connect MCP client %s: %s", config.Name, gateErr.GetErrorString()) + } + + tools := make(map[string]schemas.ChatTool) + toolNameMapping := make(map[string]string) + if externalClient == nil { + // Plugin short-circuited the connect with a success response; no live transport + // to query. Register the client as "connected" with an empty tool set — this is + // the documented Connect-success-shortcircuit gotcha. Subsequent tool calls will + // fail until a real connect happens. + m.logger.Warn("%s [%s] Connect plugin short-circuited with success; no live transport — registering with empty tool set", MCPLogPrefix, config.Name) + if connectionInfo == nil { + connectionInfo = &schemas.MCPClientConnectionInfo{Type: config.ConnectionType} + } + } else { + // Retrieve tools from the external server through the list_tools plugin gate. + // Use a bounded timeout context to prevent indefinite hangs during tool retrieval. + // For STDIO/SSE, ctx is longLivedCtx (no timeout), so we create a separate one here. + m.logger.Debug("%s [%s] Retrieving tools...", MCPLogPrefix, config.Name) + toolRetrievalCtx, toolRetrievalCancel := context.WithTimeout(m.ctx, MCPClientConnectionEstablishTimeout) + defer toolRetrievalCancel() + t, mapping, err := m.runListToolsWithHooks(toolRetrievalCtx, externalClient, config.Name) + if err != nil { + m.logger.Warn("%s Failed to retrieve tools from %s: %v", MCPLogPrefix, config.Name, err) + // Continue with connection even if tool retrieval fails + } else { + tools = t + toolNameMapping = mapping + } + m.logger.Debug("%s [%s] Retrieved %d tools", MCPLogPrefix, config.Name, len(tools)) } - m.logger.Debug("%s [%s] Retrieved %d tools", MCPLogPrefix, config.Name, len(tools)) // Second lock: Update client with final connection details and tools m.mu.Lock() @@ -1089,38 +1237,66 @@ func (m *MCPManager) connectToMCPClient(config *schemas.MCPClientConfig) error { } // createHTTPConnection creates an HTTP-based MCP client connection without holding locks. -func (m *MCPManager) createHTTPConnection(ctx context.Context, config *schemas.MCPClientConfig) (*client.Client, *schemas.MCPClientConnectionInfo, error) { +// If overrides is non-nil and carries a populated ConnectionString or Headers, those values +// are used instead of resolving them from config. This is how plugin PreHook mutations flow +// into the transport. +func (m *MCPManager) createHTTPConnection(ctx context.Context, config *schemas.MCPClientConfig, overrides *schemas.BifrostMCPConnectRequest) (*client.Client, *schemas.MCPClientConnectionInfo, error) { if config.ConnectionString == nil { return nil, nil, fmt.Errorf("HTTP connection string is required") } - // Prepare connection info - connectionInfo := &schemas.MCPClientConnectionInfo{ - Type: config.ConnectionType, - ConnectionURL: config.ConnectionString.GetValuePtr(), + + // Resolve URL (override wins) + url := config.ConnectionString.GetValue() + if overrides != nil && overrides.ConnectionString != nil { + url = *overrides.ConnectionString } - headers, err := config.HttpHeaders(ctx, m.oauth2Provider) - if err != nil { - return nil, nil, fmt.Errorf("failed to get HTTP headers: %w", err) + + // Resolve headers (override wins) + var headers map[string]string + if overrides != nil && overrides.Headers != nil { + headers = overrides.Headers + } else { + h, err := config.HttpHeaders(ctx, m.oauth2Provider) + if err != nil { + return nil, nil, fmt.Errorf("failed to get HTTP headers: %w", err) + } + headers = h } + // Create StreamableHTTP transport - httpTransport, err := transport.NewStreamableHTTP(config.ConnectionString.GetValue(), transport.WithHTTPHeaders(headers)) + httpTransport, err := transport.NewStreamableHTTP(url, transport.WithHTTPHeaders(headers)) if err != nil { return nil, nil, fmt.Errorf("failed to create HTTP transport: %w", err) } - client := client.NewClient(httpTransport) - return client, connectionInfo, nil + connectionInfo := &schemas.MCPClientConnectionInfo{ + Type: config.ConnectionType, + ConnectionURL: &url, + } + return client.NewClient(httpTransport), connectionInfo, nil } // createSTDIOConnection creates a STDIO-based MCP client connection without holding locks. -func (m *MCPManager) createSTDIOConnection(_ context.Context, config *schemas.MCPClientConfig) (*client.Client, *schemas.MCPClientConnectionInfo, error) { +// If overrides is non-nil with a populated StdioCommand/StdioArgs, those replace the config values. +func (m *MCPManager) createSTDIOConnection(_ context.Context, config *schemas.MCPClientConfig, overrides *schemas.BifrostMCPConnectRequest) (*client.Client, *schemas.MCPClientConnectionInfo, error) { if config.StdioConfig == nil { return nil, nil, fmt.Errorf("stdio config is required") } - // Prepare STDIO command info for display - cmdString := fmt.Sprintf("%s %s", config.StdioConfig.Command, strings.Join(config.StdioConfig.Args, " ")) + // Resolve command and args (override wins) + cmd := config.StdioConfig.Command + args := config.StdioConfig.Args + if overrides != nil { + if overrides.StdioCommand != nil { + cmd = *overrides.StdioCommand + } + if overrides.StdioArgs != nil { + args = overrides.StdioArgs + } + } + + cmdString := fmt.Sprintf("%s %s", cmd, strings.Join(args, " ")) - // Check if environment variables are set + // Check if environment variables are set (envs are not plugin-mutable) for _, env := range config.StdioConfig.Envs { if os.Getenv(env) == "" { return nil, nil, fmt.Errorf("environment variable %s is not set for MCP client %s", env, config.Name) @@ -1128,11 +1304,7 @@ func (m *MCPManager) createSTDIOConnection(_ context.Context, config *schemas.MC } // Create STDIO transport - stdioTransport := transport.NewStdio( - config.StdioConfig.Command, - config.StdioConfig.Envs, - config.StdioConfig.Args..., - ) + stdioTransport := transport.NewStdio(cmd, config.StdioConfig.Envs, args...) // Prepare connection info connectionInfo := &schemas.MCPClientConnectionInfo{ @@ -1140,38 +1312,43 @@ func (m *MCPManager) createSTDIOConnection(_ context.Context, config *schemas.MC StdioCommandString: &cmdString, } - client := client.NewClient(stdioTransport) - // Return nil for cmd since mark3labs/mcp-go manages the process internally - return client, connectionInfo, nil + return client.NewClient(stdioTransport), connectionInfo, nil } // createSSEConnection creates a SSE-based MCP client connection without holding locks. -func (m *MCPManager) createSSEConnection(ctx context.Context, config *schemas.MCPClientConfig) (*client.Client, *schemas.MCPClientConnectionInfo, error) { +// Same override semantics as createHTTPConnection. +func (m *MCPManager) createSSEConnection(ctx context.Context, config *schemas.MCPClientConfig, overrides *schemas.BifrostMCPConnectRequest) (*client.Client, *schemas.MCPClientConnectionInfo, error) { if config.ConnectionString == nil { return nil, nil, fmt.Errorf("SSE connection string is required") } - // Prepare connection info - connectionInfo := &schemas.MCPClientConnectionInfo{ - Type: config.ConnectionType, - ConnectionURL: config.ConnectionString.GetValuePtr(), // Reuse HTTPConnectionURL field for SSE URL display + url := config.ConnectionString.GetValue() + if overrides != nil && overrides.ConnectionString != nil { + url = *overrides.ConnectionString } - headers, err := config.HttpHeaders(ctx, m.oauth2Provider) - if err != nil { - return nil, nil, fmt.Errorf("failed to get HTTP headers: %w", err) + var headers map[string]string + if overrides != nil && overrides.Headers != nil { + headers = overrides.Headers + } else { + h, err := config.HttpHeaders(ctx, m.oauth2Provider) + if err != nil { + return nil, nil, fmt.Errorf("failed to get HTTP headers: %w", err) + } + headers = h } - // Create SSE transport - sseTransport, err := transport.NewSSE(config.ConnectionString.GetValue(), transport.WithHeaders(headers)) + sseTransport, err := transport.NewSSE(url, transport.WithHeaders(headers)) if err != nil { return nil, nil, fmt.Errorf("failed to create SSE transport: %w", err) } - client := client.NewClient(sseTransport) - - return client, connectionInfo, nil + connectionInfo := &schemas.MCPClientConnectionInfo{ + Type: config.ConnectionType, + ConnectionURL: &url, + } + return client.NewClient(sseTransport), connectionInfo, nil } // createInProcessConnection creates an in-process MCP client connection without holding locks. diff --git a/core/mcp/codemode.go b/core/mcp/codemode.go index a641289be07..81c4447015d 100644 --- a/core/mcp/codemode.go +++ b/core/mcp/codemode.go @@ -61,12 +61,6 @@ type CodeModeDependencies struct { // ClientManager provides access to MCP clients and their tools ClientManager ClientManager - // PluginPipelineProvider returns a plugin pipeline for running MCP hooks - PluginPipelineProvider func() PluginPipeline - - // ReleasePluginPipeline releases a plugin pipeline back to the pool - ReleasePluginPipeline func(pipeline PluginPipeline) - // FetchNewRequestIDFunc generates unique request IDs for nested tool calls FetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string diff --git a/core/mcp/codemode/starlark/executecode.go b/core/mcp/codemode/starlark/executecode.go index b5e5834e7fb..7769b398c71 100644 --- a/core/mcp/codemode/starlark/executecode.go +++ b/core/mcp/codemode/starlark/executecode.go @@ -467,21 +467,14 @@ func (s *StarlarkCodeMode) callMCPTool(ctx *schemas.BifrostContext, clientName, ChatAssistantMessageToolCall: &toolCallReq, } - // Check if plugin pipeline is available - if s.pluginPipelineProvider == nil { - // Should never happen, but just in case - s.logger.Warn("%s Plugin pipeline provider is nil", codemcp.CodeModeLogPrefix) - return nil, fmt.Errorf("plugin pipeline provider is nil") - } - // Get plugin pipeline and run hooks - pipeline := s.pluginPipelineProvider() + pipeline := s.clientManager.GetPluginPipeline() if pipeline == nil { // Should never happen, but just in case s.logger.Warn("%s Plugin pipeline is nil", codemcp.CodeModeLogPrefix) return nil, fmt.Errorf("plugin pipeline is nil") } - defer s.releasePluginPipeline(pipeline) + defer s.clientManager.ReleasePluginPipeline(pipeline) // Run PreMCPHooks preReq, shortCircuit, preCount := pipeline.RunMCPPreHooks(nestedCtx, mcpRequest) diff --git a/core/mcp/codemode/starlark/starlark.go b/core/mcp/codemode/starlark/starlark.go index 8b74c3fb07b..ed73bf3fe49 100644 --- a/core/mcp/codemode/starlark/starlark.go +++ b/core/mcp/codemode/starlark/starlark.go @@ -23,11 +23,9 @@ type StarlarkCodeMode struct { toolExecutionTimeout atomic.Value // time.Duration // Dependencies - clientManager mcp.ClientManager - pluginPipelineProvider func() mcp.PluginPipeline - releasePluginPipeline func(pipeline mcp.PluginPipeline) - fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string - oauth2Provider schemas.OAuth2Provider + clientManager mcp.ClientManager + fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string + oauth2Provider schemas.OAuth2Provider // Logger for this instance logger schemas.Logger @@ -84,8 +82,6 @@ func NewStarlarkCodeMode(config *mcp.CodeModeConfig, logger schemas.Logger) *Sta func (s *StarlarkCodeMode) SetDependencies(deps *mcp.CodeModeDependencies) { if deps != nil { s.clientManager = deps.ClientManager - s.pluginPipelineProvider = deps.PluginPipelineProvider - s.releasePluginPipeline = deps.ReleasePluginPipeline s.fetchNewRequestIDFunc = deps.FetchNewRequestIDFunc s.oauth2Provider = deps.OAuth2Provider } diff --git a/core/mcp/codemode/starlark/starlark_test.go b/core/mcp/codemode/starlark/starlark_test.go index a48e77a8879..3d483d95e77 100644 --- a/core/mcp/codemode/starlark/starlark_test.go +++ b/core/mcp/codemode/starlark/starlark_test.go @@ -39,6 +39,12 @@ func (m *testClientManager) GetToolPerClient(ctx context.Context) map[string][]s return m.tools } +func (m *testClientManager) GetPluginPipeline() codemcp.PluginPipeline { + return nil +} + +func (m *testClientManager) ReleasePluginPipeline(pipeline codemcp.PluginPipeline) {} + func TestStarlarkToGo(t *testing.T) { t.Run("Convert None", func(t *testing.T) { result := starlarkToGo(starlark.None) diff --git a/core/mcp/exec.go b/core/mcp/exec.go new file mode 100644 index 00000000000..a0ae952a85a --- /dev/null +++ b/core/mcp/exec.go @@ -0,0 +1,200 @@ +package mcp + +import ( + "errors" + "fmt" + "strings" + "sync" + + "github.com/maximhq/bifrost/core/schemas" +) + +// ============================================================================ +// MCP REQUEST POOL +// ============================================================================ +// +// Pool for BifrostMCPRequest objects. Owned by the mcp package because these +// requests are only used inside this package — the Bifrost public API just +// delegates to MCPManager's Execute* methods. + +var mcpRequestPool = sync.Pool{ + New: func() any { + return &schemas.BifrostMCPRequest{} + }, +} + +// resetMCPRequest zeroes a BifrostMCPRequest for reuse. Must be kept in sync with +// the fields defined on the request struct. +func resetMCPRequest(req *schemas.BifrostMCPRequest) { + req.RequestType = "" + req.ClientName = "" + req.BifrostMCPPingRequest = nil + req.BifrostMCPListToolsRequest = nil + req.BifrostMCPExecuteToolRequest = nil + req.ChatAssistantMessageToolCall = nil + req.ResponsesToolMessage = nil +} + +func getMCPRequest() *schemas.BifrostMCPRequest { + return mcpRequestPool.Get().(*schemas.BifrostMCPRequest) +} + +func releaseMCPRequest(req *schemas.BifrostMCPRequest) { + resetMCPRequest(req) + mcpRequestPool.Put(req) +} + +// ============================================================================ +// EXECUTE-TOOL GATE (matches the pattern of connect/ping/list_tools gates) +// ============================================================================ + +// executeToolWithHooks runs an MCP tool call through the plugin gate. It is the +// execute-tool counterpart to the connect/ping/list_tools gates. Mirrors the +// short-circuit + PostHook semantics of all other gates by delegating to +// runWithPluginPipeline, then adds two execute-specific touches on the returned BifrostError: +// +// - stamps ExtraFields.RequestType from the caller-provided RequestType +// - preserves MCPUserOAuthRequiredError so agent-mode detection still works +// +// requestType is the bifrost-side RequestType (ChatCompletionRequest / ResponsesRequest) +// that error metadata should carry — it isn't the same as request.RequestType. +func (m *MCPManager) executeToolWithHooks( + ctx *schemas.BifrostContext, + request *schemas.BifrostMCPRequest, + requestType schemas.RequestType, +) (*schemas.BifrostMCPResponse, *schemas.BifrostError) { + // Populate top-level ClientName from the prefixed tool name so the gate can + // attribute short-circuit responses without depending on prefix parsing. + if request != nil && request.ClientName == "" { + if toolName := request.GetToolName(); toolName != "" { + if idx := strings.IndexByte(toolName, '-'); idx > 0 { + request.ClientName = toolName[:idx] + } + } + } + + // Capture MCPUserOAuthRequiredError out-of-band: runWithPluginPipeline wraps Go errors into + // a generic BifrostError before PostHooks, which strips typed-error info. + var oauthErr *schemas.MCPUserOAuthRequiredError + + resp, bErr := m.runWithPluginPipeline(ctx, request, func(preReq *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { + result, opErr := m.ExecuteToolCall(ctx, preReq) + if opErr != nil { + errors.As(opErr, &oauthErr) + return nil, opErr + } + if result == nil { + return nil, fmt.Errorf("tool execution returned nil result") + } + return result, nil + }) + + if bErr != nil { + bErr.ExtraFields.RequestType = requestType + if oauthErr != nil { + bErr.ExtraFields.MCPAuthRequired = oauthErr + } + return nil, bErr + } + return resp, nil +} + +// executeToolForAgent is the agent-mode-facing helper. The agent loop expects a +// plain (response, error) signature and doesn't need rich BifrostError fields, +// so we collapse them. MCPUserOAuthRequiredError is returned directly when present +// so agent mode can detect it via errors.As. +func (m *MCPManager) executeToolForAgent(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { + if ctx == nil { + return nil, fmt.Errorf("context cannot be nil") + } + if request == nil { + return nil, fmt.Errorf("request cannot be nil") + } + + // Derive bifrost RequestType from the MCP request type (only execute-tool variants + // are valid in the agent loop). + var requestType schemas.RequestType + switch request.RequestType { + case schemas.MCPRequestTypeChatToolCall: + requestType = schemas.ChatCompletionRequest + case schemas.MCPRequestTypeResponsesToolCall: + requestType = schemas.ResponsesRequest + default: + return nil, fmt.Errorf("unsupported MCP request type for agent: %s", request.RequestType) + } + + resp, bErr := m.executeToolWithHooks(ctx, request, requestType) + if bErr != nil { + // Surface the typed OAuth error so agent mode can react to it. + if bErr.ExtraFields.MCPAuthRequired != nil { + return nil, bErr.ExtraFields.MCPAuthRequired + } + return nil, fmt.Errorf("tool execution failed: %s", bErr.GetErrorString()) + } + return resp, nil +} + +// ============================================================================ +// PUBLIC EXECUTE-TOOL ENTRY POINTS +// ============================================================================ + +// ExecuteChatTool executes an MCP tool call and returns the result as a chat message. +// This is the canonical entry point for manual MCP tool execution in Chat format. +// Bifrost.ExecuteChatMCPTool delegates here. +func (m *MCPManager) ExecuteChatTool(ctx *schemas.BifrostContext, toolCall *schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, *schemas.BifrostError) { + if toolCall == nil { + return nil, &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "toolCall cannot be nil"}, + ExtraFields: schemas.BifrostErrorExtraFields{RequestType: schemas.ChatCompletionRequest}, + } + } + + mcpRequest := getMCPRequest() + mcpRequest.RequestType = schemas.MCPRequestTypeChatToolCall + mcpRequest.ChatAssistantMessageToolCall = toolCall + defer releaseMCPRequest(mcpRequest) + + result, bErr := m.executeToolWithHooks(ctx, mcpRequest, schemas.ChatCompletionRequest) + if bErr != nil { + return nil, bErr + } + if result == nil || result.ChatMessage == nil { + return nil, &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "MCP tool execution returned nil chat message"}, + ExtraFields: schemas.BifrostErrorExtraFields{RequestType: schemas.ChatCompletionRequest}, + } + } + return result.ChatMessage, nil +} + +// ExecuteResponsesTool executes an MCP tool call and returns the result as a responses +// message. Bifrost.ExecuteResponsesMCPTool delegates here. +func (m *MCPManager) ExecuteResponsesTool(ctx *schemas.BifrostContext, toolCall *schemas.ResponsesToolMessage) (*schemas.ResponsesMessage, *schemas.BifrostError) { + if toolCall == nil { + return nil, &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "toolCall cannot be nil"}, + ExtraFields: schemas.BifrostErrorExtraFields{RequestType: schemas.ResponsesRequest}, + } + } + + mcpRequest := getMCPRequest() + mcpRequest.RequestType = schemas.MCPRequestTypeResponsesToolCall + mcpRequest.ResponsesToolMessage = toolCall + defer releaseMCPRequest(mcpRequest) + + result, bErr := m.executeToolWithHooks(ctx, mcpRequest, schemas.ResponsesRequest) + if bErr != nil { + return nil, bErr + } + if result == nil || result.ResponsesMessage == nil { + return nil, &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "MCP tool execution returned nil responses message"}, + ExtraFields: schemas.BifrostErrorExtraFields{RequestType: schemas.ResponsesRequest}, + } + } + return result.ResponsesMessage, nil +} diff --git a/core/mcp/healthmonitor.go b/core/mcp/healthmonitor.go index cce8ab76a5f..2f09fabff31 100644 --- a/core/mcp/healthmonitor.go +++ b/core/mcp/healthmonitor.go @@ -7,7 +7,6 @@ import ( "time" "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" "github.com/maximhq/bifrost/core/schemas" ) @@ -147,9 +146,13 @@ func (chm *ClientHealthMonitor) performHealthCheck() { clientState, exists := chm.manager.clientMap[chm.clientID] var isDisabled bool var conn *client.Client + var clientName string if exists && clientState != nil { conn = clientState.Conn isDisabled = clientState.State == schemas.MCPConnectionStateDisabled + if clientState.ExecutionConfig != nil { + clientName = clientState.ExecutionConfig.Name + } } chm.manager.mu.RUnlock() @@ -175,16 +178,13 @@ func (chm *ClientHealthMonitor) performHealthCheck() { defer cancel() if chm.isPingAvailable { - err = conn.Ping(ctx) + err = chm.runPingWithHooks(ctx, conn, clientName) } else { - listRequest := mcp.ListToolsRequest{ - PaginatedRequest: mcp.PaginatedRequest{ - Request: mcp.Request{ - Method: string(mcp.MethodToolsList), - }, - }, - } - _, err = conn.ListTools(ctx, listRequest) + // Health-check fallback uses list_tools as a liveness probe when the server + // doesn't support ping. The plugin gate fires (plugins can observe / mutate / + // short-circuit) but the resulting tools are DISCARDED — periodic tool sync + // owns tool state, this path is liveness-only. + _, _, err = chm.manager.runListToolsWithHooks(ctx, conn, clientName) } } diff --git a/core/mcp/interface.go b/core/mcp/interface.go index 316e138bc0b..9442661d06f 100644 --- a/core/mcp/interface.go +++ b/core/mcp/interface.go @@ -28,24 +28,31 @@ type MCPManagerInterface interface { UpdateToolManagerConfig(config *schemas.MCPToolManagerConfig) // Agent Mode Operations - // CheckAndExecuteAgentForChatRequest handles agent mode for Chat Completions API + // CheckAndExecuteAgentForChatRequest handles agent mode for Chat Completions API. + // Tool executions inside the agent loop go through the plugin gate internally — + // callers no longer inject an executeTool function. CheckAndExecuteAgentForChatRequest( ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest, response *schemas.BifrostChatResponse, makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError), - executeTool func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error), ) (*schemas.BifrostChatResponse, *schemas.BifrostError) - // CheckAndExecuteAgentForResponsesRequest handles agent mode for Responses API + // CheckAndExecuteAgentForResponsesRequest handles agent mode for Responses API. + // Tool executions inside the agent loop go through the plugin gate internally. CheckAndExecuteAgentForResponsesRequest( ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest, response *schemas.BifrostResponsesResponse, makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError), - executeTool func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error), ) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) + // ExecuteChatTool / ExecuteResponsesTool run a single MCP tool call through the + // plugin gate and return the result in the appropriate API format. Bifrost's + // ExecuteChatMCPTool / ExecuteResponsesMCPTool delegate here. + ExecuteChatTool(ctx *schemas.BifrostContext, toolCall *schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, *schemas.BifrostError) + ExecuteResponsesTool(ctx *schemas.BifrostContext, toolCall *schemas.ResponsesToolMessage) (*schemas.ResponsesMessage, *schemas.BifrostError) + // Client Management // GetClients returns all MCP clients GetClients() []schemas.MCPClientState diff --git a/core/mcp/mcp.go b/core/mcp/mcp.go index 9a3fba0938f..57fa99f9854 100644 --- a/core/mcp/mcp.go +++ b/core/mcp/mcp.go @@ -42,6 +42,12 @@ type MCPManager struct { healthMonitorManager *HealthMonitorManager // Manager for client health monitors toolSyncManager *ToolSyncManager // Manager for periodic tool synchronization reconnectingClients sync.Map // Tracks in-flight reconnect attempts per client ID (map[string]bool) + + // Plugin pipeline access for connect/ping/list_tools hooks. nil-safe — gates short-circuit + // to the underlying op when no pipeline is configured. Also used by ToolsManager for the + // existing execute-tool hooks. + pluginPipelineProvider func() PluginPipeline + releasePluginPipeline func(pipeline PluginPipeline) } // MCPToolFunction is a generic function type for handling tool calls with typed arguments. @@ -103,7 +109,9 @@ func NewMCPManager(ctx context.Context, config schemas.MCPConfig, oauth2Provider } } - manager.toolsManager = NewToolsManager(config.ToolManagerConfig, manager, config.FetchNewRequestIDFunc, pluginPipelineProvider, releasePluginPipeline, oauth2Provider, logger) + manager.pluginPipelineProvider = pluginPipelineProvider + manager.releasePluginPipeline = releasePluginPipeline + manager.toolsManager = NewToolsManager(config.ToolManagerConfig, manager, config.FetchNewRequestIDFunc, oauth2Provider, logger) // Set up CodeMode if provided - inject dependencies after manager is created if codeMode != nil { @@ -161,7 +169,23 @@ func NewMCPManager(ctx context.Context, config schemas.MCPConfig, oauth2Provider // ToolsManager and CodeMode. Call this after attaching an externally-created MCPManager to a Bifrost // instance so that nested tool calls in code mode can run through Bifrost's plugin hooks. func (manager *MCPManager) SetPluginPipeline(provider func() PluginPipeline, release func(PluginPipeline)) { - manager.toolsManager.SetPluginPipeline(provider, release) + manager.pluginPipelineProvider = provider + manager.releasePluginPipeline = release +} + +// GetPluginPipeline returns a plugin pipeline from the provider, or nil if no provider is configured. +func (manager *MCPManager) GetPluginPipeline() PluginPipeline { + if manager.pluginPipelineProvider != nil { + return manager.pluginPipelineProvider() + } + return nil +} + +// ReleasePluginPipeline releases a plugin pipeline back to the pool via the configured release function. +func (manager *MCPManager) ReleasePluginPipeline(pipeline PluginPipeline) { + if manager.releasePluginPipeline != nil { + manager.releasePluginPipeline(pipeline) + } } // AddToolsToRequest parses available MCP tools from the context and adds them to the request. @@ -235,7 +259,6 @@ func (m *MCPManager) CheckAndExecuteAgentForChatRequest( req *schemas.BifrostChatRequest, response *schemas.BifrostChatResponse, makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError), - executeTool func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error), ) (*schemas.BifrostChatResponse, *schemas.BifrostError) { if makeReq == nil { return nil, &schemas.BifrostError{ @@ -250,8 +273,9 @@ func (m *MCPManager) CheckAndExecuteAgentForChatRequest( m.logger.Debug("No tool calls detected, returning response") return response, nil } - // Execute agent mode - return m.toolsManager.ExecuteAgentForChatRequest(ctx, req, response, makeReq, executeTool) + // Execute agent mode. The agent's tool executions go through the plugin gate + // internally via m.executeToolForAgent — no external callback injection needed. + return m.toolsManager.ExecuteAgentForChatRequest(ctx, req, response, makeReq, m.executeToolForAgent) } // CheckAndExecuteAgentForResponsesRequest checks if the responses response contains tool calls, @@ -287,7 +311,6 @@ func (m *MCPManager) CheckAndExecuteAgentForResponsesRequest( req *schemas.BifrostResponsesRequest, response *schemas.BifrostResponsesResponse, makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError), - executeTool func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error), ) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) { if makeReq == nil { return nil, &schemas.BifrostError{ @@ -302,8 +325,8 @@ func (m *MCPManager) CheckAndExecuteAgentForResponsesRequest( m.logger.Debug("No tool calls detected, returning response") return response, nil } - // Execute agent mode - return m.toolsManager.ExecuteAgentForResponsesRequest(ctx, req, response, makeReq, executeTool) + // Execute agent mode. Tool executions go through the plugin gate internally. + return m.toolsManager.ExecuteAgentForResponsesRequest(ctx, req, response, makeReq, m.executeToolForAgent) } // Cleanup performs cleanup of all MCP resources including clients and local server. diff --git a/core/mcp/pluginpipeline.go b/core/mcp/pluginpipeline.go new file mode 100644 index 00000000000..34b2e1b88d5 --- /dev/null +++ b/core/mcp/pluginpipeline.go @@ -0,0 +1,378 @@ +package mcp + +import ( + "context" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/mark3labs/mcp-go/client" + "github.com/maximhq/bifrost/core/schemas" +) + +// MCPOpFunc is the closure each call site provides to runWithPluginPipeline. It receives the +// (possibly mutated) request that flowed through PreHooks and is responsible for +// performing the wire call (including any internal retries) and building a +// BifrostMCPResponse from the outcome. The plain Go error returned here is wrapped +// into a BifrostError by the gate before being handed to PostMCPHooks. +type MCPOpFunc func(preReq *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) + +// runWithPluginPipeline wraps an MCP wire operation (connect / ping / list_tools / execute_tool) +// with the plugin pipeline. It is the single source of truth for the MCP plugin gate +// pattern — handleMCPToolExecution in core/bifrost.go calls into this same function. +// +// 1. Acquire pipeline (no-op pass-through if none configured) +// 2. Run PreMCPHooks — plugins may mutate the request or short-circuit +// 3. On short-circuit: invoke PostMCPHooks with the short-circuit outcome, +// drain plugin logs, return +// 4. Otherwise: invoke op with the mutated request, then run PostMCPHooks on +// the outcome (response or error), drain plugin logs, return +// +// The op closure is responsible for reading mutated values from preReq's per-op +// sub-request struct (Headers, ConnectionString, ChatAssistantMessageToolCall, etc.) +// and using them for the actual wire call. +// +// Returns *BifrostError so callers can preserve rich error fields (AllowFallbacks, +// MCPAuthRequired). +func (m *MCPManager) runWithPluginPipeline( + ctx *schemas.BifrostContext, + req *schemas.BifrostMCPRequest, + op MCPOpFunc, +) (*schemas.BifrostMCPResponse, *schemas.BifrostError) { + // Ensure a request ID exists so plugin hooks have something to correlate on. + // Connect/ping/list_tools fire from background contexts that typically lack one. + if ctx != nil { + if _, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string); !ok { + ctx.SetValue(schemas.BifrostContextKeyRequestID, uuid.New().String()) + } + } + + // Wrap the whole gate (PreHook + op + PostHook) in an outer span so traces show one + // row per MCP op alongside the per-plugin spans the pipeline emits internally. + tracer, _ := ctx.Value(schemas.BifrostContextKeyTracer).(schemas.Tracer) + var spanHandle schemas.SpanHandle + if tracer != nil { + spanName := fmt.Sprintf("mcp.%s", req.RequestType) + if req.ClientName != "" { + spanName = fmt.Sprintf("%s.%s", spanName, req.ClientName) + } + _, spanHandle = tracer.StartSpan(ctx, spanName, schemas.SpanKindMCPClient) + } + defer func() { + if tracer != nil { + tracer.EndSpan(spanHandle, schemas.SpanStatusOk, "") + } + }() + + // MCP request type stamped on every wrapped BifrostError so downstream gates + // (governance, logging) can discriminate execute-tool calls from ping/list_tools. + mcpReqType := schemas.MCPRequestType("") + if req != nil { + mcpReqType = req.RequestType + } + + // No pipeline configured → run the op directly, no hooks (but span still recorded). + pipeline := m.GetPluginPipeline() + if pipeline == nil { + resp, opErr := op(req) + if opErr != nil { + return resp, &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: opErr.Error()}, + ExtraFields: schemas.BifrostErrorExtraFields{MCPRequestType: mcpReqType}, + } + } + return resp, nil + } + defer m.ReleasePluginPipeline(pipeline) + + // PreHooks. preReq is the (possibly mutated) request that must flow to the op. + preReq, shortCircuit, preCount := pipeline.RunMCPPreHooks(ctx, req) + + // Pull attribution from the request so short-circuit responses still carry + // ClientName/ToolName when PostHooks or observers consume them. + clientName := "" + toolName := "" + if preReq != nil { + clientName = preReq.ClientName + toolName = preReq.GetToolName() + } else if req != nil { + clientName = req.ClientName + toolName = req.GetToolName() + } + + if shortCircuit != nil { + // Short-circuit with response — still run PostHooks for plugins that ran. + if shortCircuit.Response != nil { + shortCircuit.Response.PopulateExtraFields(mcpReqType, clientName, toolName) + finalResp, finalErr := pipeline.RunMCPPostHooks(ctx, shortCircuit.Response, nil, preCount) + drainMCPPluginLogs(ctx) + if finalErr != nil { + return nil, finalErr + } + return finalResp, nil + } + // Short-circuit with error — still run PostHooks (they may recover). + if shortCircuit.Error != nil { + if shortCircuit.Error.ExtraFields.MCPRequestType == "" { + shortCircuit.Error.ExtraFields.MCPRequestType = mcpReqType + } + finalResp, finalErr := pipeline.RunMCPPostHooks(ctx, nil, shortCircuit.Error, preCount) + drainMCPPluginLogs(ctx) + if finalErr != nil { + return nil, finalErr + } + if finalResp != nil { + finalResp.PopulateExtraFields(mcpReqType, clientName, toolName) + return finalResp, nil + } + return nil, shortCircuit.Error + } + } + + if preReq == nil { + return nil, &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "MCP request after plugin hooks cannot be nil"}, + ExtraFields: schemas.BifrostErrorExtraFields{MCPRequestType: mcpReqType}, + } + } + + // Run the actual wire op with the mutated request. + resp, opErr := op(preReq) + if resp != nil { + resp.PopulateExtraFields(mcpReqType, clientName, toolName) + } + + // Wrap opErr as BifrostError so PostHooks see a typed error. + var bErr *schemas.BifrostError + if opErr != nil { + bErr = &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: opErr.Error()}, + ExtraFields: schemas.BifrostErrorExtraFields{MCPRequestType: mcpReqType}, + } + } + + finalResp, finalErr := pipeline.RunMCPPostHooks(ctx, resp, bErr, preCount) + drainMCPPluginLogs(ctx) + + if finalErr != nil { + return finalResp, finalErr + } + return finalResp, nil +} + +// MCPConnectOpFunc is the closure each Connect call site provides to +// +// . It receives the (possibly mutated) typed sub-request +// +// that flowed through PreMCPConnectionHook plugins and performs the actual transport +// + initialize work (with internal retries), returning a typed sub-response. +type MCPConnectOpFunc func(preReq *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectResponse, error) + +// runConnectWithPluginPipeline is the typed Connect-specific counterpart to +// runWithPluginPipeline. Connect ops bypass the envelope-based pipeline entirely: +// plugins implement MCPConnectionPlugin (not MCPPlugin), the request/response types +// are the typed sub-structs, and the dispatch never wraps anything in +// BifrostMCPRequest/BifrostMCPResponse. +func (m *MCPManager) runConnectWithPluginPipeline( + ctx *schemas.BifrostContext, + req *schemas.BifrostMCPConnectRequest, + op MCPConnectOpFunc, +) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError) { + if ctx != nil { + if _, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string); !ok { + ctx.SetValue(schemas.BifrostContextKeyRequestID, uuid.New().String()) + } + } + + clientName := "" + if req != nil { + clientName = req.ClientName + } + + // Outer span so traces show one row per Connect op. + tracer, _ := ctx.Value(schemas.BifrostContextKeyTracer).(schemas.Tracer) + var spanHandle schemas.SpanHandle + if tracer != nil { + spanName := "mcp.connect" + if clientName != "" { + spanName = fmt.Sprintf("%s.%s", spanName, clientName) + } + _, spanHandle = tracer.StartSpan(ctx, spanName, schemas.SpanKindMCPClient) + } + defer func() { + if tracer != nil { + tracer.EndSpan(spanHandle, schemas.SpanStatusOk, "") + } + }() + + // No pipeline configured → run the op directly. + pipeline := m.GetPluginPipeline() + if pipeline == nil { + resp, opErr := op(req) + if opErr != nil { + return resp, &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: opErr.Error()}, + } + } + return resp, nil + } + defer m.ReleasePluginPipeline(pipeline) + + preReq, shortCircuit, preCount := pipeline.RunMCPPreConnectionHooks(ctx, req) + if preReq != nil { + clientName = preReq.ClientName + } + + if shortCircuit != nil { + if shortCircuit.Response != nil { + shortCircuit.Response.PopulateExtraFields(clientName) + finalResp, finalErr := pipeline.RunMCPPostConnectionHooks(ctx, shortCircuit.Response, nil, preCount) + drainMCPPluginLogs(ctx) + if finalErr != nil { + return nil, finalErr + } + return finalResp, nil + } + if shortCircuit.Error != nil { + finalResp, finalErr := pipeline.RunMCPPostConnectionHooks(ctx, nil, shortCircuit.Error, preCount) + drainMCPPluginLogs(ctx) + if finalErr != nil { + return nil, finalErr + } + if finalResp != nil { + finalResp.PopulateExtraFields(clientName) + return finalResp, nil + } + return nil, shortCircuit.Error + } + } + + if preReq == nil { + return nil, &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: "Connect request after plugin hooks cannot be nil"}, + } + } + + resp, opErr := op(preReq) + if resp != nil { + resp.PopulateExtraFields(clientName) + } + + var bErr *schemas.BifrostError + if opErr != nil { + bErr = &schemas.BifrostError{ + IsBifrostError: false, + Error: &schemas.ErrorField{Message: opErr.Error()}, + } + } + + finalResp, finalErr := pipeline.RunMCPPostConnectionHooks(ctx, resp, bErr, preCount) + drainMCPPluginLogs(ctx) + + if finalErr != nil { + return finalResp, finalErr + } + return finalResp, nil +} + +// drainMCPPluginLogs mirrors bifrost.drainAndAttachPluginLogs for the mcp package. +// It attaches accumulated plugin log entries to the active trace, if any. +func drainMCPPluginLogs(ctx *schemas.BifrostContext) { + if ctx == nil { + return + } + tracer, _ := ctx.Value(schemas.BifrostContextKeyTracer).(schemas.Tracer) + if tracer == nil { + return + } + traceID, _ := ctx.Value(schemas.BifrostContextKeyTraceID).(string) + if traceID == "" { + return + } + logs := ctx.DrainPluginLogs() + if len(logs) == 0 { + return + } + tracer.AttachPluginLogs(traceID, logs) +} + +// runListToolsWithHooks wraps retrieveExternalToolsDetailed in the MCP plugin gate. +// All four list_tools call sites (connect / oauth-verify / sync / health-check fallback) +// go through this helper so plugins see one consistent hook. The PostHook may mutate +// the Tools / ToolNameMapping fields on the response — the caller receives the mutated +// values via the returned maps. +// +// A PreHook short-circuit with Response is treated as a synthetic tool list (used as-is). +// A PreHook short-circuit with Error returns the error; the caller decides whether to +// keep existing state. +func (m *MCPManager) runListToolsWithHooks(ctx context.Context, conn *client.Client, clientName string) (map[string]schemas.ChatTool, map[string]string, error) { + req := &schemas.BifrostMCPRequest{ + RequestType: schemas.MCPRequestTypeListTools, + ClientName: clientName, + BifrostMCPListToolsRequest: &schemas.BifrostMCPListToolsRequest{}, + } + gateCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) + start := time.Now() + + resp, bErr := m.runWithPluginPipeline(gateCtx, req, func(preReq *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { + detailed, opErr := retrieveExternalToolsDetailed(ctx, conn, clientName, m.logger) + if opErr != nil { + return nil, opErr + } + return &schemas.BifrostMCPResponse{ + BifrostMCPListToolsResponse: &schemas.BifrostMCPListToolsResponse{ + Tools: detailed.tools, + ToolNameMapping: detailed.toolNameMapping, + RawToolCount: detailed.rawCount, + SkippedTools: detailed.skipped, + }, + ExtraFields: schemas.BifrostMCPResponseExtraFields{ + Latency: time.Since(start).Milliseconds(), + }, + }, nil + }) + + if bErr != nil { + return nil, nil, fmt.Errorf("failed to list tools: %s", bErr.GetErrorString()) + } + if resp == nil || resp.BifrostMCPListToolsResponse == nil { + // Defensive: response somehow lost its list_tools payload (e.g. PostHook nilled it). + // Surface empty maps rather than nil to mirror the underlying list_tools contract. + return make(map[string]schemas.ChatTool), make(map[string]string), nil + } + return resp.Tools, resp.ToolNameMapping, nil +} + +// runPingWithHooks wraps conn.Ping in the MCP plugin gate. A PreHook may short-circuit +// the ping (synthetic healthy/unhealthy) without touching the wire; a PostHook may +// inspect the latency and outcome. Any error returned here is treated identically to a +// real ping failure by the health-monitor state machine. +func (chm *ClientHealthMonitor) runPingWithHooks(ctx context.Context, conn *client.Client, clientName string) error { + req := &schemas.BifrostMCPRequest{ + RequestType: schemas.MCPRequestTypePing, + ClientName: clientName, + BifrostMCPPingRequest: &schemas.BifrostMCPPingRequest{}, + } + gateCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) + start := time.Now() + _, bErr := chm.manager.runWithPluginPipeline(gateCtx, req, func(preReq *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) { + if pingErr := conn.Ping(ctx); pingErr != nil { + return nil, pingErr + } + return &schemas.BifrostMCPResponse{ + BifrostMCPPingResponse: &schemas.BifrostMCPPingResponse{}, + ExtraFields: schemas.BifrostMCPResponseExtraFields{ + Latency: time.Since(start).Milliseconds(), + }, + }, nil + }) + if bErr != nil { + return fmt.Errorf("ping failed: %s", bErr.GetErrorString()) + } + return nil +} diff --git a/core/mcp/toolmanager.go b/core/mcp/toolmanager.go index 3f8bb2f11bd..6cbfee5ea49 100644 --- a/core/mcp/toolmanager.go +++ b/core/mcp/toolmanager.go @@ -22,13 +22,22 @@ type ClientManager interface { GetClientByName(clientName string) *schemas.MCPClientState GetClientForTool(toolName string) *schemas.MCPClientState GetToolPerClient(ctx context.Context) map[string][]schemas.ChatTool + GetPluginPipeline() PluginPipeline + ReleasePluginPipeline(pipeline PluginPipeline) } // PluginPipeline represents the plugin execution pipeline interface -// This allows ToolsManager to run plugin hooks without direct dependency on Bifrost +// This allows ToolsManager to run plugin hooks without direct dependency on Bifrost. +// Two parallel pipelines exist: the envelope-based MCP pipeline for Ping/ListTools/ +// ExecuteTool variants, and the typed Connect pipeline for MCPConnectionPlugin. type PluginPipeline interface { + // Envelope pipeline (Ping / ListTools / ExecuteTool variants) RunMCPPreHooks(ctx *schemas.BifrostContext, req *schemas.BifrostMCPRequest) (*schemas.BifrostMCPRequest, *schemas.MCPPluginShortCircuit, int) RunMCPPostHooks(ctx *schemas.BifrostContext, mcpResp *schemas.BifrostMCPResponse, bifrostErr *schemas.BifrostError, runFrom int) (*schemas.BifrostMCPResponse, *schemas.BifrostError) + + // Typed Connect pipeline (MCPConnectionPlugin) + RunMCPPreConnectionHooks(ctx *schemas.BifrostContext, req *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectRequest, *schemas.MCPConnectionShortCircuit, int) + RunMCPPostConnectionHooks(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPConnectResponse, bifrostErr *schemas.BifrostError, runFrom int) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError) } // ToolsManager manages MCP tool execution and agent mode. @@ -51,13 +60,6 @@ type ToolsManager struct { // This id is attached to ctx.Value(schemas.BifrostContextKeyRequestID) in the agent mode. // If not provided, same request ID is used for all tool call result messages without any overrides. fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string - - // Function to get a plugin pipeline from the pool for running MCP plugin hooks - // Used when executeCode tool calls nested MCP tools to ensure plugins run for them - pluginPipelineProvider func() PluginPipeline - - // Function to release a plugin pipeline back to the pool - releasePluginPipeline func(pipeline PluginPipeline) } // NewToolsManager creates and initializes a new tools manager instance. @@ -68,8 +70,6 @@ type ToolsManager struct { // - config: Tool manager configuration with execution timeout and max agent depth // - clientManager: Client manager interface for accessing MCP clients and tools // - fetchNewRequestIDFunc: Optional function to generate unique request IDs for agent mode -// - pluginPipelineProvider: Optional function to get a plugin pipeline for running MCP hooks -// - releasePluginPipeline: Optional function to release a plugin pipeline back to the pool // // Returns: // - *ToolsManager: Initialized tools manager instance @@ -77,8 +77,6 @@ func NewToolsManager( config *schemas.MCPToolManagerConfig, clientManager ClientManager, fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string, - pluginPipelineProvider func() PluginPipeline, - releasePluginPipeline func(pipeline PluginPipeline), oauth2Provider schemas.OAuth2Provider, logger schemas.Logger, ) *ToolsManager { @@ -86,8 +84,6 @@ func NewToolsManager( config, clientManager, fetchNewRequestIDFunc, - pluginPipelineProvider, - releasePluginPipeline, nil, // Use default code mode (will be set later via SetCodeMode) oauth2Provider, logger, @@ -101,8 +97,6 @@ func NewToolsManager( // - config: Tool manager configuration with execution timeout and max agent depth // - clientManager: Client manager interface for accessing MCP clients and tools // - fetchNewRequestIDFunc: Optional function to generate unique request IDs for agent mode -// - pluginPipelineProvider: Optional function to get a plugin pipeline for running MCP hooks -// - releasePluginPipeline: Optional function to release a plugin pipeline back to the pool // - codeMode: Optional CodeMode implementation (if nil, must be set later via SetCodeMode) // // Returns: @@ -111,8 +105,6 @@ func NewToolsManagerWithCodeMode( config *schemas.MCPToolManagerConfig, clientManager ClientManager, fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string, - pluginPipelineProvider func() PluginPipeline, - releasePluginPipeline func(pipeline PluginPipeline), codeMode CodeMode, oauth2Provider schemas.OAuth2Provider, logger schemas.Logger, @@ -144,14 +136,12 @@ func NewToolsManagerWithCodeMode( } manager := &ToolsManager{ - clientManager: clientManager, - fetchNewRequestIDFunc: fetchNewRequestIDFunc, - pluginPipelineProvider: pluginPipelineProvider, - releasePluginPipeline: releasePluginPipeline, - codeMode: codeMode, - logger: logger, - agentModeExecutor: agentModeExecutor, - oauth2Provider: oauth2Provider, + clientManager: clientManager, + fetchNewRequestIDFunc: fetchNewRequestIDFunc, + codeMode: codeMode, + logger: logger, + agentModeExecutor: agentModeExecutor, + oauth2Provider: oauth2Provider, } // Initialize atomic values @@ -178,23 +168,9 @@ func (m *ToolsManager) GetCodeMode() CodeMode { // This is useful when constructing a CodeMode implementation externally. func (m *ToolsManager) GetCodeModeDependencies() *CodeModeDependencies { return &CodeModeDependencies{ - ClientManager: m.clientManager, - PluginPipelineProvider: m.pluginPipelineProvider, - ReleasePluginPipeline: m.releasePluginPipeline, - FetchNewRequestIDFunc: m.fetchNewRequestIDFunc, - OAuth2Provider: m.oauth2Provider, - } -} - -// SetPluginPipeline updates the plugin pipeline provider and release function -// on both the ToolsManager and its CodeMode implementation. -// This is used when an externally-created MCPManager is attached to a Bifrost instance -// via SetMCPManager, so the CodeMode can route nested tool calls through Bifrost's plugin hooks. -func (m *ToolsManager) SetPluginPipeline(provider func() PluginPipeline, release func(PluginPipeline)) { - m.pluginPipelineProvider = provider - m.releasePluginPipeline = release - if m.codeMode != nil { - m.codeMode.SetDependencies(m.GetCodeModeDependencies()) + ClientManager: m.clientManager, + FetchNewRequestIDFunc: m.fetchNewRequestIDFunc, + OAuth2Provider: m.oauth2Provider, } } @@ -734,18 +710,9 @@ func (m *ToolsManager) executeToolInternal(ctx *schemas.BifrostContext, toolCall } // ExecuteAgentForChatRequest executes agent mode for a chat request, handling -// iterative tool calls up to the configured maximum depth. It delegates to the -// shared agent execution logic with the manager's configuration and dependencies. -// -// Parameters: -// - ctx: Context for agent execution -// - req: The original chat request -// - resp: The initial chat response containing tool calls -// - makeReq: Function to make subsequent chat requests during agent execution -// -// Returns: -// - *schemas.BifrostChatResponse: The final response after agent execution -// - *schemas.BifrostError: Any error that occurred during agent execution +// iterative tool calls up to the configured maximum depth. Tool executions inside +// the agent loop are dispatched through the executeTool callback the caller provides +// (typically MCPManager.executeToolForAgent, which routes through the plugin gate). func (m *ToolsManager) ExecuteAgentForChatRequest( ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest, @@ -753,10 +720,10 @@ func (m *ToolsManager) ExecuteAgentForChatRequest( makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError), executeTool func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error), ) (*schemas.BifrostChatResponse, *schemas.BifrostError) { - // Use provided executeTool function, or fall back to internal ExecuteTool - executeToolFunc := executeTool - if executeToolFunc == nil { - executeToolFunc = m.ExecuteTool + // Defensive: if no executor was supplied, fall back to the un-hooked path. This + // path is only exercised by internal callers that have already gated above. + if executeTool == nil { + executeTool = m.ExecuteTool } return m.agentModeExecutor.ExecuteAgentForChatRequest( ctx, @@ -765,24 +732,12 @@ func (m *ToolsManager) ExecuteAgentForChatRequest( resp, makeReq, m.fetchNewRequestIDFunc, - executeToolFunc, + executeTool, m.clientManager, ) } -// ExecuteAgentForResponsesRequest executes agent mode for a responses request, handling -// iterative tool calls up to the configured maximum depth. It delegates to the -// shared agent execution logic with the manager's configuration and dependencies. -// -// Parameters: -// - ctx: Context for agent execution -// - req: The original responses request -// - resp: The initial responses response containing tool calls -// - makeReq: Function to make subsequent responses requests during agent execution -// -// Returns: -// - *schemas.BifrostResponsesResponse: The final response after agent execution -// - *schemas.BifrostError: Any error that occurred during agent execution +// ExecuteAgentForResponsesRequest mirrors ExecuteAgentForChatRequest for the Responses API. func (m *ToolsManager) ExecuteAgentForResponsesRequest( ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest, @@ -790,10 +745,8 @@ func (m *ToolsManager) ExecuteAgentForResponsesRequest( makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError), executeTool func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error), ) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) { - // Use provided executeTool function, or fall back to internal ExecuteTool - executeToolFunc := executeTool - if executeToolFunc == nil { - executeToolFunc = m.ExecuteTool + if executeTool == nil { + executeTool = m.ExecuteTool } return m.agentModeExecutor.ExecuteAgentForResponsesRequest( ctx, @@ -802,7 +755,7 @@ func (m *ToolsManager) ExecuteAgentForResponsesRequest( resp, makeReq, m.fetchNewRequestIDFunc, - executeToolFunc, + executeTool, m.clientManager, ) } @@ -903,7 +856,6 @@ func ExecuteToolWithUserToken(ctx context.Context, config *schemas.MCPClientConf return tempClient.CallTool(ctx, callRequest) } - // GetCodeModeBindingLevel returns the current code mode binding level. // This method is safe to call concurrently from multiple goroutines. func (m *ToolsManager) GetCodeModeBindingLevel() schemas.CodeModeBindingLevel { diff --git a/core/mcp/toolmanager_test.go b/core/mcp/toolmanager_test.go index 27751b2b3d8..0ee7d042209 100644 --- a/core/mcp/toolmanager_test.go +++ b/core/mcp/toolmanager_test.go @@ -22,10 +22,10 @@ func (m *mockToolClientManager) GetClientByName(clientName string) *schemas.MCPC return &schemas.MCPClientState{ Name: "test-client", ExecutionConfig: &schemas.MCPClientConfig{ - ID: "test-client", - Name: "test-client", + ID: "test-client", + Name: "test-client", IsCodeModeClient: false, - ToolsToExecute: []string{"*"}, + ToolsToExecute: []string{"*"}, }, } } @@ -42,6 +42,9 @@ func (m *mockToolClientManager) GetToolPerClient(ctx context.Context) map[string } } +func (m *mockToolClientManager) GetPluginPipeline() PluginPipeline { return nil } +func (m *mockToolClientManager) ReleasePluginPipeline(pipeline PluginPipeline) {} + // makeTool is a convenience constructor for test tool fixtures. func makeTool(name string) schemas.ChatTool { return schemas.ChatTool{ @@ -102,8 +105,6 @@ func newToolsManagerForTest(cm ClientManager) *ToolsManager { }, cm, nil, // fetchNewRequestIDFunc - nil, // pluginPipelineProvider - nil, // releasePluginPipeline nil, // oauth2Provider &MockLogger{}, ) @@ -154,8 +155,8 @@ func TestBuildIntegrationDuplicateCheckMap_NilFunction_IsSkipped(t *testing.T) { tools := []schemas.ChatTool{ {Type: schemas.ChatToolTypeFunction, Function: nil}, // nil Function - makeTool(""), // empty name - makeTool("valid_tool"), // valid + makeTool(""), // empty name + makeTool("valid_tool"), // valid } m := buildIntegrationDuplicateCheckMap(tools, "", defaultLogger) diff --git a/core/mcp/toolsync.go b/core/mcp/toolsync.go index 042df1f2cf0..3a9ffda4c5d 100644 --- a/core/mcp/toolsync.go +++ b/core/mcp/toolsync.go @@ -130,7 +130,7 @@ func (cts *ClientToolSyncer) performSync() { ctx, cancel := context.WithTimeout(context.Background(), cts.timeout) defer cancel() - newTools, newMapping, err := retrieveExternalTools(ctx, conn, clientName, cts.logger) + newTools, newMapping, err := cts.manager.runListToolsWithHooks(ctx, conn, clientName) if err != nil { // On failure, keep existing tools intact cts.logger.Warn("%s Tool sync failed for %s, keeping existing tools: %v", MCPLogPrefix, cts.clientID, err) diff --git a/core/mcp/utils.go b/core/mcp/utils.go index d261c1c26ea..173efe51286 100644 --- a/core/mcp/utils.go +++ b/core/mcp/utils.go @@ -288,10 +288,21 @@ func ExecuteWithRetry( return lastErr } -// retrieveExternalTools retrieves and filters tools from an external MCP server without holding locks. -// Uses exponential backoff retry logic (5 retries, 1-30 seconds) for tool retrieval. -// Returns both the tools map and a name mapping (sanitized_name -> original_mcp_name) for tool execution. -func retrieveExternalTools(ctx context.Context, client *client.Client, clientName string, logger schemas.Logger) (map[string]schemas.ChatTool, map[string]string, error) { +// listToolsResult captures the full outcome of retrieveExternalToolsDetailed so the +// plugin gate can expose RawToolCount and SkippedTools to plugins. +type listToolsResult struct { + tools map[string]schemas.ChatTool + toolNameMapping map[string]string + rawCount int + skipped []schemas.SkippedMCPTool +} + +// retrieveExternalToolsDetailed retrieves and filters tools from an external MCP server +// without holding locks. Uses exponential backoff retry logic (5 retries, 1-30 seconds) +// for tool retrieval. Returns the full result so the plugin gate can surface RawToolCount +// and SkippedTools. All callers should go through MCPManager.runListToolsWithHooks rather +// than calling this directly — the gate wraps this with PreMCPHook/PostMCPHook. +func retrieveExternalToolsDetailed(ctx context.Context, client *client.Client, clientName string, logger schemas.Logger) (*listToolsResult, error) { // Get available tools from external server with retry logic listRequest := mcp.ListToolsRequest{ PaginatedRequest: mcp.PaginatedRequest{ @@ -314,15 +325,19 @@ func retrieveExternalTools(ctx context.Context, client *client.Client, clientNam logger, ) if err != nil { - return nil, nil, fmt.Errorf("failed to list tools after %d retries: %v", retryConfig.MaxRetries, err) + return nil, fmt.Errorf("failed to list tools after %d retries: %v", retryConfig.MaxRetries, err) } if toolsResponse == nil { - return make(map[string]schemas.ChatTool), make(map[string]string), nil // No tools available + return &listToolsResult{ + tools: make(map[string]schemas.ChatTool), + toolNameMapping: make(map[string]string), + }, nil } tools := make(map[string]schemas.ChatTool) toolNameMapping := make(map[string]string) // Maps sanitized_name -> original_mcp_name + var skipped []schemas.SkippedMCPTool // toolsResponse is already a ListToolsResult for _, mcpTool := range toolsResponse.Tools { @@ -330,6 +345,10 @@ func retrieveExternalTools(ctx context.Context, client *client.Client, clientNam validationName := strings.ReplaceAll(mcpTool.Name, "-", "_") if err := validateNormalizedToolName(validationName); err != nil { logger.Warn("%s Skipping MCP tool %q: %v", MCPLogPrefix, mcpTool.Name, err) + skipped = append(skipped, schemas.SkippedMCPTool{ + OriginalName: mcpTool.Name, + Reason: err.Error(), + }) continue } @@ -349,7 +368,12 @@ func retrieveExternalTools(ctx context.Context, client *client.Client, clientNam toolNameMapping[sanitizedToolName] = mcpTool.Name } - return tools, toolNameMapping, nil + return &listToolsResult{ + tools: tools, + toolNameMapping: toolNameMapping, + rawCount: len(toolsResponse.Tools), + skipped: skipped, + }, nil } // shouldIncludeClient determines if a client should be included based on filtering rules. diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index cd9a6101f7c..d5410d74188 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -766,19 +766,91 @@ func (br *BifrostRequest) SetRawRequestBody(rawRequestBody []byte) { type MCPRequestType string const ( + MCPRequestTypePing MCPRequestType = "ping" + MCPRequestTypeListTools MCPRequestType = "list_tools" + + // [DEPRECATED] these will be replaced by MCPRequestTypeExecuteTool in the next major bump, but are kept for backward compatibility for now since some tools still rely on the old fields MCPRequestTypeChatToolCall MCPRequestType = "chat_tool_call" // Chat API format MCPRequestTypeResponsesToolCall MCPRequestType = "responses_tool_call" // Responses API format + + // Will be used in from the next major bump + MCPRequestTypeExecuteTool MCPRequestType = "execute_tool" ) -// BifrostMCPRequest is the request struct for all MCP requests. -// only ONE of the following fields should be set: -// - ChatAssistantMessageToolCall -// - ResponsesToolMessage +// IsExecuteTool reports whether this is one of the execute-tool request variants +// (Chat, Responses, or the future unified ExecuteTool). Used by MCPPlugin pre/post +// hooks to skip non-tool envelope ops (Ping/ListTools) without sniffing pointer +// fields on the request/response. +// +// NOTE: this helper exists because three execute-tool request types currently +// coexist for backwards compat (ChatToolCall + ResponsesToolCall are deprecated). +// Once callers fully migrate to MCPRequestTypeExecuteTool, this method will be +// removed and consumers should switch to `t == MCPRequestTypeExecuteTool` directly. +func (t MCPRequestType) IsExecuteTool() bool { + switch t { + case MCPRequestTypeChatToolCall, + MCPRequestTypeResponsesToolCall, + MCPRequestTypeExecuteTool: + return true + } + return false +} + +// BifrostMCPRequest is the envelope for MCP requests that flow through the generic +// PreMCPHook/PostMCPHook pipeline (Ping, ListTools, ExecuteTool variants). Connect +// requests do NOT use this envelope — they are dispatched via the typed +// MCPConnectionPlugin interface using *BifrostMCPConnectRequest directly. +// +// Exactly one of the embedded sub-request pointers is populated, matched by RequestType: +// - RequestType == MCPRequestTypePing → BifrostMCPPingRequest +// - RequestType == MCPRequestTypeListTools → BifrostMCPListToolsRequest +// - RequestType == MCPRequestTypeExecuteTool / MCPRequestTypeChatToolCall / MCPRequestTypeResponsesToolCall → BifrostMCPExecuteToolRequest type BifrostMCPRequest struct { RequestType MCPRequestType + ClientName string // MCP client this request targets (always set, regardless of request type) + *BifrostMCPPingRequest + *BifrostMCPListToolsRequest + + // [DEPRECATED] these will be replaced by BifrostMCPExecuteToolRequest in the next major bump, but are kept for backward compatibility for now since some tools still rely on the old fields *ChatAssistantMessageToolCall *ResponsesToolMessage + + // Will be used in from the next major bump + *BifrostMCPExecuteToolRequest +} + +// BifrostMCPConnectRequest carries the prepared inputs for an MCP connect operation. +// Fields marked "mutable" may be modified by a plugin's PreMCPHook and the mutated values +// will be used for the actual transport creation; "observe-only" fields are passed to plugins +// for context but mutations are ignored (changing the transport type mid-flight would break +// the rest of the connect codepath). +type BifrostMCPConnectRequest struct { + ClientName string // observe-only — name of the client being connected + ConnectionType MCPConnectionType // observe-only — transport type being established (http/stdio/sse/inprocess) + AuthType MCPAuthType // observe-only — authentication mode configured on the client + ConnectionString *string // mutable — URL for http/sse, nil for stdio/inprocess + Headers map[string]string // mutable — transport-level headers (http/sse only; stdio/inprocess ignore) + StdioCommand *string // mutable — command for stdio connections (nil otherwise) + StdioArgs []string // mutable — argv for stdio connections (nil otherwise) +} + +// BifrostMCPPingRequest is intentionally empty: the wire ping rides over the existing +// transport and has no per-call headers or parameters. Plugins observe via ClientName on +// the parent BifrostMCPRequest and may short-circuit (synthetic healthy/unhealthy). +type BifrostMCPPingRequest struct { +} + +// BifrostMCPListToolsRequest is intentionally empty for the same reason as ping: list_tools +// reuses the existing transport's headers. Plugins observe via ClientName and may short-circuit +// (e.g. cached tool list). +type BifrostMCPListToolsRequest struct { +} + +// Keeping the stub for now, will be used from the next major bump when we remove the old ChatToolCall and ResponsesToolMessage fields. +// Note that the tool name and arguments are not standardized in this struct yet since they are still being pulled from the old fields for backward compatibility, +// but they will be standardized in the future when we remove the old fields. +type BifrostMCPExecuteToolRequest struct { } func (r *BifrostMCPRequest) GetToolName() string { @@ -1193,14 +1265,109 @@ func (r *BifrostResponse) PopulateExtraFields(requestType RequestType, provider } } -// BifrostMCPResponse is the response struct for all MCP responses. -// only ONE of the following fields should be set: -// - ChatMessage -// - ResponsesMessage +// BifrostMCPResponse is the envelope for MCP responses that flow through the generic +// PostMCPHook pipeline (Ping, ListTools, ExecuteTool variants). Connect responses do +// NOT use this envelope — they are dispatched via the typed MCPConnectionPlugin +// interface using *BifrostMCPConnectResponse directly. +// +// Exactly one of the embedded sub-response pointers is populated, matched by the +// originating request's RequestType. ExtraFields (ClientName / ToolName / Latency) +// applies to all envelope variants. For execute-tool requests in the back-compat +// window, the direct ChatMessage / ResponsesMessage fields are populated instead of +// any embedded sub-response. type BifrostMCPResponse struct { + *BifrostMCPPingResponse + *BifrostMCPListToolsResponse + + // [DEPRECATED] back-compat fields for execute-tool requests; will move into + // BifrostMCPExecuteToolResponse in the next major bump. ChatMessage *ChatMessage ResponsesMessage *ResponsesMessage - ExtraFields BifrostMCPResponseExtraFields + + // Empty stub today; will hold ChatMessage/ResponsesMessage in the next major bump. + *BifrostMCPExecuteToolResponse + + ExtraFields BifrostMCPResponseExtraFields +} + +// Latency for envelope MCP responses (ping, list_tools, execute_tool) is reported via +// BifrostMCPResponse.ExtraFields.Latency (milliseconds). Connect carries its own +// ExtraFields on BifrostMCPConnectResponse below — see typed Connect path. + +type BifrostMCPConnectResponse struct { + ConnectionInfo *MCPClientConnectionInfo // Connection metadata after the handshake completes + ServerInfo *MCPServerInfo // Name + version from the initialize handshake + ProtocolVersion string // Negotiated MCP protocol version + ServerCapabilities *MCPServerCapabilities // Which MCP feature groups the server claims to support + ExtraFields BifrostMCPResponseExtraFields +} + +// PopulateExtraFields backfills ClientName on the Connect response when it's not +// already set. Mirrors BifrostMCPResponse.PopulateExtraFields. Connect has no tool +// name, so only ClientName is populated. +func (r *BifrostMCPConnectResponse) PopulateExtraFields(clientName string) { + if r == nil { + return + } + if r.ExtraFields.ClientName == "" { + r.ExtraFields.ClientName = clientName + } +} + +// MCPServerInfo mirrors the ServerInfo portion of the MCP initialize handshake. +type MCPServerInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// MCPServerCapabilities mirrors the high-level capability flags from the MCP initialize handshake. +// Only the booleans Bifrost cares about today; can grow as needed. +type MCPServerCapabilities struct { + Tools bool `json:"tools"` // server supports tools/list + tools/call + Resources bool `json:"resources"` // server supports resources + Prompts bool `json:"prompts"` // server supports prompts + Logging bool `json:"logging"` // server supports logging +} + +type BifrostMCPPingResponse struct { +} + +type BifrostMCPListToolsResponse struct { + Tools map[string]ChatTool // Discovered tools keyed by client-prefixed name + ToolNameMapping map[string]string // sanitized_name -> original_mcp_name + RawToolCount int // Count returned by the MCP server before Bifrost-side filtering + SkippedTools []SkippedMCPTool // Tools Bifrost dropped during conversion + reason +} + +// SkippedMCPTool describes a tool that the MCP server returned but Bifrost did not include +// in the final tool map (e.g. invalid normalized name). +type SkippedMCPTool struct { + OriginalName string `json:"original_name"` + Reason string `json:"reason"` +} + +// Keeping the stub for now, will be used from the next major bump when we move ChatMessage +// and ResponsesMessage into this struct. +type BifrostMCPExecuteToolResponse struct { +} + +// PopulateExtraFields backfills ExtraFields.{MCPRequestType, ClientName, ToolName} +// when they aren't already set on the response. Mirrors BifrostResponse.PopulateExtraFields +// and is used by every MCP gate to ensure short-circuit responses carry the same +// attribution as real wire-call responses. +func (r *BifrostMCPResponse) PopulateExtraFields(mcpRequestType MCPRequestType, clientName, toolName string) { + if r == nil { + return + } + if r.ExtraFields.MCPRequestType == "" { + r.ExtraFields.MCPRequestType = mcpRequestType + } + if r.ExtraFields.ClientName == "" { + r.ExtraFields.ClientName = clientName + } + if r.ExtraFields.ToolName == "" { + r.ExtraFields.ToolName = toolName + } } // BifrostResponseExtraFields contains additional fields in a response. @@ -1221,9 +1388,10 @@ type BifrostResponseExtraFields struct { } type BifrostMCPResponseExtraFields struct { - ClientName string `json:"client_name"` - ToolName string `json:"tool_name"` - Latency int64 `json:"latency"` // in milliseconds + MCPRequestType MCPRequestType `json:"mcp_request_type"` // request type this response corresponds to — lets PostMCPHook discriminate ping/list_tools from tool execute on success too + ClientName string `json:"client_name"` + ToolName string `json:"tool_name"` // empty for all but MCPRequestTypeExecuteTool requests for backwards compat, will be a pointer from next major bump. + Latency int64 `json:"latency"` // in milliseconds } // BifrostCacheDebug represents debug information about the cache. @@ -1337,6 +1505,45 @@ func (e *BifrostError) String() string { return string(b) } +func (e *BifrostError) GetErrorString() string { + if e == nil { + return "" + } + if e.Error != nil && e.Error.Message != "" { + return e.Error.Message + } else if e.StatusCode != nil { + switch *e.StatusCode { + case 401: + return "unauthorized" + case 403: + return "forbidden" + case 404: + return "endpoint not found" + case 405: + return "method not allowed" + case 429: + return "rate limit exceeded" + case 500: + return "internal server error" + case 502: + return "bad gateway" + case 503: + return "service unavailable" + case 504: + return "gateway timeout" + default: + if e.Error != nil && e.Error.Message != "" { + return e.Error.Message + } + return fmt.Sprintf("HTTP %d error", *e.StatusCode) + } + } else if e.Type != nil { + return *e.Type + } else { + return "unknown error" + } +} + // StreamControl represents stream control options. type StreamControl struct { LogError *bool `json:"log_error,omitempty"` // Optional: Controls logging of error @@ -1414,6 +1621,7 @@ type BifrostErrorExtraFields struct { OriginalModelRequested string `json:"original_model_requested,omitempty"` ResolvedModelUsed string `json:"resolved_model_used,omitempty"` RequestType RequestType `json:"request_type,omitempty"` + MCPRequestType MCPRequestType `json:"mcp_request_type,omitempty"` RawRequest interface{} `json:"raw_request,omitempty"` RawResponse interface{} `json:"raw_response,omitempty"` ConvertedRequestType RequestType `json:"converted_request_type,omitempty"` diff --git a/core/schemas/plugin.go b/core/schemas/plugin.go index 5e0d0687182..d526390a523 100644 --- a/core/schemas/plugin.go +++ b/core/schemas/plugin.go @@ -268,6 +268,47 @@ type MCPPlugin interface { PostMCPHook(ctx *BifrostContext, resp *BifrostMCPResponse, bifrostErr *BifrostError) (*BifrostMCPResponse, *BifrostError, error) } +// MCPConnectionPlugin is an optional, typed extension interface for handling MCP +// Connect events. Connect is morally separate from the other MCP lifecycle ops +// (Ping/ListTools/ExecuteTool) — it establishes the transport before a usable +// client exists, and carries transport-level inputs (URL, headers, stdio args) +// that don't apply post-connection. Plugins implementing this interface receive +// Connect events via the typed methods; their generic PreMCPHook/PostMCPHook +// (if also implemented) is NOT called for Connect requests. +// +// Plugins registered via MCPPlugins must still satisfy MCPPlugin. To write a +// plugin that only handles Connect events, embed MCPPluginNoOpHooks for free +// no-op implementations of the generic Pre/PostMCPHook. +// +// NOTE (backwards compat): keeping the Connect hooks on a separate optional +// interface — and the MCPPluginNoOpHooks helper — is purely a backwards-compat +// shim so existing MCPPlugin implementations don't break with the addition of +// Connect hooks. In a future major release these two methods will move onto +// MCPPlugin directly and every MCP plugin will be required to implement them. +type MCPConnectionPlugin interface { + MCPPlugin + + PreMCPConnectionHook(ctx *BifrostContext, req *BifrostMCPConnectRequest) (*BifrostMCPConnectRequest, *MCPConnectionShortCircuit, error) + PostMCPConnectionHook(ctx *BifrostContext, resp *BifrostMCPConnectResponse, bifrostErr *BifrostError) (*BifrostMCPConnectResponse, *BifrostError, error) +} + +// MCPPluginNoOpHooks provides no-op implementations of PreMCPHook and PostMCPHook. +// Embed this in plugins that only want to implement an extension interface +// (e.g. MCPConnectionPlugin) and don't need to observe the generic hook surface. +// +// The plugin must still provide its own GetName and Cleanup (from BasePlugin). +type MCPPluginNoOpHooks struct{} + +// PreMCPHook returns the request unchanged with no short-circuit. +func (MCPPluginNoOpHooks) PreMCPHook(_ *BifrostContext, req *BifrostMCPRequest) (*BifrostMCPRequest, *MCPPluginShortCircuit, error) { + return req, nil, nil +} + +// PostMCPHook returns the response and error unchanged. +func (MCPPluginNoOpHooks) PostMCPHook(_ *BifrostContext, resp *BifrostMCPResponse, bifrostErr *BifrostError) (*BifrostMCPResponse, *BifrostError, error) { + return resp, bifrostErr, nil +} + // Plugin placement constants control where custom plugins execute relative to built-in plugins. type PluginPlacement string diff --git a/core/schemas/plugin_native.go b/core/schemas/plugin_native.go index c8d2a46852b..4256a62aa9d 100644 --- a/core/schemas/plugin_native.go +++ b/core/schemas/plugin_native.go @@ -31,6 +31,14 @@ type MCPPluginShortCircuit struct { Error *BifrostError // If set, short-circuit with this error (can set AllowFallbacks field) } +// MCPConnectionShortCircuit is the typed short-circuit for MCPConnectionPlugin. +// It carries a typed Connect response (instead of the generic envelope) so plugin +// authors don't have to wrap responses in BifrostMCPResponse. +type MCPConnectionShortCircuit struct { + Response *BifrostMCPConnectResponse // If set, short-circuit with this synthetic Connect outcome + Error *BifrostError // If set, short-circuit with this error +} + // PluginShortCircuit is the legacy name for LLMPluginShortCircuit (v1.3.x compatibility). // Deprecated: Use LLMPluginShortCircuit instead. type PluginShortCircuit = LLMPluginShortCircuit diff --git a/core/schemas/trace.go b/core/schemas/trace.go index 227af9122cb..916285b8b49 100644 --- a/core/schemas/trace.go +++ b/core/schemas/trace.go @@ -164,6 +164,9 @@ const ( SpanKindPlugin SpanKind = "plugin" // SpanKindMCPTool represents an MCP tool invocation SpanKindMCPTool SpanKind = "mcp.tool" + // SpanKindMCPClient represents an MCP client lifecycle operation (connect/ping/list_tools). + // These run in the background per-client and are not part of an LLM request flow. + SpanKindMCPClient SpanKind = "mcp.client" // SpanKindRetry represents a retry attempt SpanKindRetry SpanKind = "retry" // SpanKindFallback represents a fallback to another provider diff --git a/core/utils.go b/core/utils.go index 16ddc126a01..424fe0d0d80 100644 --- a/core/utils.go +++ b/core/utils.go @@ -409,43 +409,9 @@ func MarshalUnsafe(v any) string { return strings.TrimSpace(buf.String()) } +// // [Deprecated] use err.GetErrorString() instead. Will be removed in a future release. func GetErrorMessage(err *schemas.BifrostError) string { - if err == nil { - return "" - } - if err.Error != nil && err.Error.Message != "" { - return err.Error.Message - } else if err.StatusCode != nil { - switch *err.StatusCode { - case 401: - return "unauthorized" - case 403: - return "forbidden" - case 404: - return "endpoint not found" - case 405: - return "method not allowed" - case 429: - return "rate limit exceeded" - case 500: - return "internal server error" - case 502: - return "bad gateway" - case 503: - return "service unavailable" - case 504: - return "gateway timeout" - default: - if err.Error != nil && err.Error.Message != "" { - return err.Error.Message - } - return fmt.Sprintf("HTTP %d error", *err.StatusCode) - } - } else if err.Type != nil { - return *err.Type - } else { - return "unknown error" - } + return err.GetErrorString() } // GetStringFromContext safely extracts a string value from context diff --git a/framework/plugins/soloader.go b/framework/plugins/soloader.go index a4fd6251a96..face5772a75 100644 --- a/framework/plugins/soloader.go +++ b/framework/plugins/soloader.go @@ -132,6 +132,23 @@ func (l *SharedObjectPluginLoader) LoadPlugin(path string, config any) (schemas. } } + // Optional: PreMCPConnectionHook (MCPConnectionPlugin — typed Connect hook). + // New .so plugins built against MCPConnectionPlugin can export this symbol to + // observe Connect events. Legacy plugins that don't export it keep working; + // DynamicPlugin's default PreMCPConnectionHook is a no-op passthrough. + if sym, err := pluginObj.Lookup("PreMCPConnectionHook"); err == nil { + if dp.preMCPConnectionHook, ok = sym.(func(ctx *schemas.BifrostContext, req *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectRequest, *schemas.MCPConnectionShortCircuit, error)); !ok { + return nil, fmt.Errorf("failed to cast PreMCPConnectionHook to expected signature") + } + } + + // Optional: PostMCPConnectionHook (MCPConnectionPlugin — typed Connect hook). + if sym, err := pluginObj.Lookup("PostMCPConnectionHook"); err == nil { + if dp.postMCPConnectionHook, ok = sym.(func(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPConnectResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError, error)); !ok { + return nil, fmt.Errorf("failed to cast PostMCPConnectionHook to expected signature") + } + } + // Optional: Inject (ObservabilityPlugin) if sym, err := pluginObj.Lookup("Inject"); err == nil { if dp.inject, ok = sym.(func(ctx context.Context, trace *schemas.Trace) error); !ok { diff --git a/framework/plugins/soplugin.go b/framework/plugins/soplugin.go index 8f0447826cb..196a1ceed27 100644 --- a/framework/plugins/soplugin.go +++ b/framework/plugins/soplugin.go @@ -34,6 +34,13 @@ type DynamicPlugin struct { preMCPHook func(ctx *schemas.BifrostContext, req *schemas.BifrostMCPRequest) (*schemas.BifrostMCPRequest, *schemas.MCPPluginShortCircuit, error) postMCPHook func(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPResponse, *schemas.BifrostError, error) + // MCPConnectionPlugin (optional, typed). Forward-compat: new .so plugins can + // export PreMCPConnectionHook/PostMCPConnectionHook to receive Connect events + // with the typed signatures. Legacy plugins (pre-MCPConnectionPlugin) leave + // these nil and silently no-op for Connect. + preMCPConnectionHook func(ctx *schemas.BifrostContext, req *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectRequest, *schemas.MCPConnectionShortCircuit, error) + postMCPConnectionHook func(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPConnectResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError, error) + // ObservabilityPlugin (optional) inject func(ctx context.Context, trace *schemas.Trace) error } @@ -104,6 +111,26 @@ func (dp *DynamicPlugin) PostMCPHook(ctx *schemas.BifrostContext, resp *schemas. return dp.postMCPHook(ctx, resp, bifrostErr) } +// PreMCPConnectionHook satisfies MCPConnectionPlugin for dynamically-loaded plugins. +// If the .so exported PreMCPConnectionHook, dispatch to it. Otherwise default to +// a no-op passthrough — legacy plugins predating MCPConnectionPlugin keep working +// as MCPPlugin (via PreMCPHook/PostMCPHook) and silently skip Connect events. +func (dp *DynamicPlugin) PreMCPConnectionHook(ctx *schemas.BifrostContext, req *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectRequest, *schemas.MCPConnectionShortCircuit, error) { + if dp.preMCPConnectionHook == nil { + return req, nil, nil + } + return dp.preMCPConnectionHook(ctx, req) +} + +// PostMCPConnectionHook satisfies MCPConnectionPlugin for dynamically-loaded plugins. +// Same dispatch as PreMCPConnectionHook: typed symbol if exported, else no-op. +func (dp *DynamicPlugin) PostMCPConnectionHook(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPConnectResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError, error) { + if dp.postMCPConnectionHook == nil { + return resp, bifrostErr, nil + } + return dp.postMCPConnectionHook(ctx, resp, bifrostErr) +} + // Inject receives completed traces for observability backends (ObservabilityPlugin interface) func (dp *DynamicPlugin) Inject(ctx context.Context, trace *schemas.Trace) error { if dp.inject == nil { diff --git a/plugins/governance/main.go b/plugins/governance/main.go index 0f9aeba1623..d87dcb0e932 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -1424,6 +1424,11 @@ func (p *GovernancePlugin) PostLLMHook(ctx *schemas.BifrostContext, result *sche func (p *GovernancePlugin) PreMCPHook(ctx *schemas.BifrostContext, req *schemas.BifrostMCPRequest) (*schemas.BifrostMCPRequest, *schemas.MCPPluginShortCircuit, error) { toolName := req.GetToolName() + // Skip for non tool execution requests + if !req.RequestType.IsExecuteTool() { + return req, nil, nil + } + // Skip governance for codemode tools if bifrost.IsCodemodeTool(toolName) { return req, nil, nil @@ -1501,6 +1506,19 @@ func (p *GovernancePlugin) PostMCPHook(ctx *schemas.BifrostContext, resp *schema return resp, bifrostErr, nil } + // Skip non tool-execute envelopes. The MCP gate stamps MCPRequestType on both + // the success response (BifrostMCPResponse.ExtraFields) and the error + // (BifrostError.ExtraFields), so a single check covers both paths. + mcpReqType := schemas.MCPRequestType("") + if resp != nil { + mcpReqType = resp.ExtraFields.MCPRequestType + } else if bifrostErr != nil { + mcpReqType = bifrostErr.ExtraFields.MCPRequestType + } + if !mcpReqType.IsExecuteTool() { + return resp, bifrostErr, nil + } + // Extract governance information virtualKey := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyVirtualKey) requestID := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyRequestID) diff --git a/plugins/logging/main.go b/plugins/logging/main.go index e9599229356..b843e290a85 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -1147,6 +1147,11 @@ func (p *LoggerPlugin) PreMCPHook(ctx *schemas.BifrostContext, req *schemas.Bifr return req, nil, nil } + // Only log for tool execute requests + if !req.RequestType.IsExecuteTool() { + return req, nil, nil + } + requestID, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string) if !ok || requestID == "" { p.logger.Error("request-id not found in context or is empty in PreMCPHook") @@ -1254,8 +1259,20 @@ func (p *LoggerPlugin) PostMCPHook(ctx *schemas.BifrostContext, resp *schemas.Bi return resp, bifrostErr, nil } + // Skip non tool-execute envelopes (Ping/ListTools). The MCP gate stamps + // MCPRequestType on both the success response and the error, so a single check + // covers both paths — no pending MCP log entry was created in PreMCPHook for + // anything but execute-tool requests. + mcpReqType := schemas.MCPRequestType("") + if resp != nil { + mcpReqType = resp.ExtraFields.MCPRequestType + } else if bifrostErr != nil { + mcpReqType = bifrostErr.ExtraFields.MCPRequestType + } + if !mcpReqType.IsExecuteTool() { + return resp, bifrostErr, nil + } // Skip logging for codemode tools (executeToolCode, listToolFiles, readToolFile) - // We check the tool name from the response instead of context flags if resp != nil && bifrost.IsCodemodeTool(resp.ExtraFields.ToolName) { return resp, bifrostErr, nil } diff --git a/ui/app/workspace/config/views/mcpView.tsx b/ui/app/workspace/config/views/mcpView.tsx index 8067241cc10..3bed1d4f1c3 100644 --- a/ui/app/workspace/config/views/mcpView.tsx +++ b/ui/app/workspace/config/views/mcpView.tsx @@ -350,7 +350,7 @@ export default function MCPView() {

External Base URLs

Override Bifrost's public base URL when it runs behind a reverse proxy. In most setups - both URLs are the same — leave them blank to derive the URL from the incoming{" "} + both URLs are the same — leave them blank to derive the URL from the incoming{" "} Host header. Both fields support env var syntax (e.g.{" "} env.BIFROST_EXTERNAL_URL).