feat: propagate request context through mcp client connection

core/bifrost.go

@@ -3641,12 +3641,12 @@ func (bifrost *Bifrost) GetAvailableMCPTools(ctx *schemas.BifrostContext) []sche
 //
 // Example:
 //
-//	err := bifrost.AddMCPClient(schemas.MCPClientConfig{
+//	err := bifrost.AddMCPClient(ctx, &schemas.MCPClientConfig{
 //	    Name: "my-mcp-client",
 //	    ConnectionType: schemas.MCPConnectionTypeHTTP,
 //	    ConnectionString: &url,
 //	})
-func (bifrost *Bifrost) AddMCPClient(config *schemas.MCPClientConfig) error {
+func (bifrost *Bifrost) AddMCPClient(ctx context.Context, config *schemas.MCPClientConfig) error {
 	if bifrost.MCPManager == nil {
 		// Use sync.Once to ensure thread-safe initialization
 		bifrost.mcpInitOnce.Do(func() {
@@ -3674,7 +3674,7 @@ func (bifrost *Bifrost) AddMCPClient(config *schemas.MCPClientConfig) error {
 		return fmt.Errorf("MCP manager is not initialized")
 	}
 
-	return bifrost.MCPManager.AddClient(config)
+	return bifrost.MCPManager.AddClient(ctx, config)
 }
 
 // RemoveMCPClient removes an MCP client from the Bifrost instance.

core/internal/mcptests/agent_filtering_test.go

@@ -1,6 +1,7 @@
 package mcptests
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -507,7 +508,7 @@ func TestAgent_FilteringWithMultipleClients(t *testing.T) {
 	tempConfig := GetTemperatureMCPClientConfig("")
 	tempConfig.ToolsToAutoExecute = []string{} // Not auto-executed
 
-	err = manager.AddClient(&tempConfig)
+	err = manager.AddClient(context.Background(), &tempConfig)
 	if err != nil {
 		t.Skipf("Skipping test - temperature server not available: %v", err)
 		return
@@ -528,7 +529,7 @@ func TestAgent_FilteringWithMultipleClients(t *testing.T) {
 					ID:   schemas.Ptr("call-2"),
 					Type: schemas.Ptr("function"),
 					Function: schemas.ChatAssistantMessageToolCallFunction{
-						Name: schemas.Ptr("bifrostInternal-get_temperature"),
+						Name:      schemas.Ptr("bifrostInternal-get_temperature"),
 						Arguments: `{"location": "New York"}`,
 					},
 				},
@@ -593,7 +594,7 @@ func TestAgent_ToolConflictInAgentMode(t *testing.T) {
 	tempConfig := GetTemperatureMCPClientConfig("")
 	tempConfig.ToolsToAutoExecute = []string{} // Not auto
 
-	err = manager.AddClient(&tempConfig)
+	err = manager.AddClient(context.Background(), &tempConfig)
 	if err != nil {
 		t.Skipf("Skipping test - temperature server not available: %v", err)
 		return
@@ -611,7 +612,7 @@ func TestAgent_ToolConflictInAgentMode(t *testing.T) {
 					ID:   schemas.Ptr("call-1"),
 					Type: schemas.Ptr("function"),
 					Function: schemas.ChatAssistantMessageToolCallFunction{
-						Name: schemas.Ptr("bifrostInternal-get_temperature"),
+						Name:      schemas.Ptr("bifrostInternal-get_temperature"),
 						Arguments: `{"location": "New York"}`,
 					},
 				},
@@ -823,7 +824,7 @@ func TestAgent_Filtering_ResponsesFormat(t *testing.T) {
 			CreateResponsesResponseWithToolCalls([]schemas.ResponsesToolMessage{
 				{
 					CallID:    schemas.Ptr("call-1"),
-					Name: schemas.Ptr("bifrostInternal-echo"),
+					Name:      schemas.Ptr("bifrostInternal-echo"),
 					Arguments: schemas.Ptr(`{"message": "responses format"}`),
 				},
 			}),

core/internal/mcptests/client_management_test.go

@@ -1,6 +1,7 @@
 package mcptests
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -25,11 +26,11 @@ func TestAddClientDuplicate(t *testing.T) {
 
 	// Add client
 	clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
-	err := manager.AddClient(&clientConfig)
+	err := manager.AddClient(context.Background(), &clientConfig)
 	require.NoError(t, err, "should add client first time")
 
 	// Try to add same client again
-	err = manager.AddClient(&clientConfig)
+	err = manager.AddClient(context.Background(), &clientConfig)
 	// Should either return error or be idempotent
 	if err == nil {
 		clients := manager.GetClients()
@@ -431,7 +432,7 @@ func TestConcurrentClientOperations(t *testing.T) {
 			clientConfig.ID = string(rune('a'+id)) + "-concurrent-client"
 			clientConfig.Name = "TestHTTPServer" + string(rune('a'+id))
 
-			err := manager.AddClient(&clientConfig)
+			err := manager.AddClient(context.Background(), &clientConfig)
 			if err != nil {
 				errors <- err
 			}

core/internal/mcptests/concurrency_advanced_test.go

@@ -162,7 +162,7 @@ func TestConcurrent_AddRemoveClients(t *testing.T) {
 				ToolsToAutoExecute: []string{},
 			}
 
-			err := manager.AddClient(&clientConfig)
+			err := manager.AddClient(context.Background(), &clientConfig)
 			if err != nil {
 				// InProcess connections without a server instance will fail
 				// This is expected - we're just testing that the operations are concurrent and don't deadlock

core/internal/mcptests/connect_ping_listtools_test.go

@@ -6,11 +6,11 @@ import (
 	"testing"
 	"time"
 
+	mcpgo "github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
 	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"
 )
@@ -91,7 +91,7 @@ func TestConnectHook_FiresOnAddClient(t *testing.T) {
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
 
 	cfg := inProcessClientConfig("connect_fires", buildInProcessServer(t))
-	require.NoError(t, manager.AddClient(cfg))
+	require.NoError(t, manager.AddClient(context.Background(), cfg))
 
 	pre := plugin.GetPreHookCalls()
 	post := plugin.GetPostHookCalls()
@@ -111,7 +111,7 @@ func TestConnectHook_PostHookPopulatesServerInfo(t *testing.T) {
 	plugin := NewTestConnectPlugin()
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
 
-	require.NoError(t, manager.AddClient(inProcessClientConfig("server_info", buildInProcessServer(t))))
+	require.NoError(t, manager.AddClient(context.Background(), inProcessClientConfig("server_info", buildInProcessServer(t))))
 
 	post := plugin.GetPostHookCalls()
 	require.Len(t, post, 1)
@@ -147,7 +147,7 @@ func TestConnectHook_PreHookShortCircuitError_FailsAddClient(t *testing.T) {
 	})
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
 
-	err := manager.AddClient(inProcessClientConfig("blocked_client", buildInProcessServer(t)))
+	err := manager.AddClient(context.Background(), inProcessClientConfig("blocked_client", buildInProcessServer(t)))
 	require.Error(t, err, "Plugin error short-circuit should fail AddClient")
 	assert.Contains(t, err.Error(), "blocked by governance")
 
@@ -176,7 +176,7 @@ func TestConnectHook_PreHookShortCircuitResponse_RegistersWithEmptyTools(t *test
 
 	// 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))))
+	require.NoError(t, manager.AddClient(context.Background(), inProcessClientConfig("synthetic_client", buildInProcessServer(t))))
 
 	clients := manager.GetClients()
 	var found *schemas.MCPClientState
@@ -223,7 +223,7 @@ func TestConnectHook_AuthorizationHidden_HeadersAuth(t *testing.T) {
 		ToolsToExecute: []string{"*"},
 	}
 	// Short-circuit returns an error → AddClient fails. That's expected.
-	_ = manager.AddClient(cfg)
+	_ = manager.AddClient(context.Background(), cfg)
 
 	calls := plugin.GetPreHookCalls()
 	require.NotEmpty(t, calls, "PreHook should have fired before short-circuit")
@@ -276,7 +276,7 @@ func TestListToolsHook_FiresOnAddClient(t *testing.T) {
 	plugin := NewTestListToolsPlugin()
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
 
-	require.NoError(t, manager.AddClient(inProcessClientConfig("list_fires", buildInProcessServer(t))))
+	require.NoError(t, manager.AddClient(context.Background(), inProcessClientConfig("list_fires", buildInProcessServer(t))))
 
 	pre := plugin.GetPreHookCalls()
 	post := plugin.GetPostHookCalls()
@@ -306,7 +306,7 @@ func TestListToolsHook_PostHookFilterAppliedToClientState(t *testing.T) {
 	})
 
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
-	require.NoError(t, manager.AddClient(inProcessClientConfig("list_filter", buildInProcessServer(t))))
+	require.NoError(t, manager.AddClient(context.Background(), 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.
@@ -345,7 +345,7 @@ func TestListToolsHook_PreHookShortCircuitWithSyntheticTools(t *testing.T) {
 	plugin.SetShortCircuitResponse(synthetic)
 
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
-	require.NoError(t, manager.AddClient(inProcessClientConfig("list_synth", buildInProcessServer(t))))
+	require.NoError(t, manager.AddClient(context.Background(), inProcessClientConfig("list_synth", buildInProcessServer(t))))
 
 	clients := manager.GetClients()
 	var target *schemas.MCPClientState
@@ -375,7 +375,7 @@ func TestListToolsHook_PreHookShortCircuitError_LeavesEmptyToolMap(t *testing.T)
 	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))))
+	require.NoError(t, manager.AddClient(context.Background(), inProcessClientConfig("list_err", buildInProcessServer(t))))
 
 	clients := manager.GetClients()
 	var target *schemas.MCPClientState
@@ -397,7 +397,7 @@ func TestListToolsHook_FiresOnConnectAndAgain(t *testing.T) {
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
 
 	cfg := inProcessClientConfig("list_reconnect", buildInProcessServer(t))
-	require.NoError(t, manager.AddClient(cfg))
+	require.NoError(t, manager.AddClient(context.Background(), 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.
@@ -416,7 +416,7 @@ func TestPingHook_FiresViaHealthMonitor(t *testing.T) {
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
 
 	cfg := inProcessClientConfig("ping_fires", buildInProcessServer(t))
-	require.NoError(t, manager.AddClient(cfg))
+	require.NoError(t, manager.AddClient(context.Background(), cfg))
 
 	// AddClient starts its own health monitor at 10s interval — far too slow for
 	// tests. Spin up a dedicated monitor at 10ms instead.
@@ -454,7 +454,7 @@ func TestPingHook_PreHookShortCircuitHealthy(t *testing.T) {
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
 
 	cfg := inProcessClientConfig("ping_healthy", buildInProcessServer(t))
-	require.NoError(t, manager.AddClient(cfg))
+	require.NoError(t, manager.AddClient(context.Background(), cfg))
 
 	monitor := mcp.NewClientHealthMonitor(manager, cfg.ID, 10*time.Millisecond, true, core.NewDefaultLogger(schemas.LogLevelError))
 	monitor.Start()
@@ -488,7 +488,7 @@ func TestPingHook_PreHookShortCircuitError_DoesNotPanic(t *testing.T) {
 
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{plugin})
 	cfg := inProcessClientConfig("ping_err", buildInProcessServer(t))
-	require.NoError(t, manager.AddClient(cfg))
+	require.NoError(t, manager.AddClient(context.Background(), cfg))
 
 	monitor := mcp.NewClientHealthMonitor(manager, cfg.ID, 10*time.Millisecond, true, core.NewDefaultLogger(schemas.LogLevelError))
 	monitor.Start()
@@ -517,7 +517,7 @@ func TestPingHook_DoesNotFireWhenPingUnavailable(t *testing.T) {
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{pingPlugin, listPlugin})
 
 	cfg := inProcessClientConfig("ping_unavailable", buildInProcessServer(t))
-	require.NoError(t, manager.AddClient(cfg))
+	require.NoError(t, manager.AddClient(context.Background(), cfg))
 
 	// Reset the list-tools plugin so we ignore the AddClient-time invocation.
 	listPlugin.Reset()
@@ -550,7 +550,7 @@ func TestMCPGate_AllRequestTypesCarryClientName(t *testing.T) {
 	logPlugin := NewTestLoggingPlugin()
 	manager, _ := setupBifrostWithPlugins(t, []schemas.MCPPlugin{logPlugin})
 
-	require.NoError(t, manager.AddClient(inProcessClientConfig("client_name_test", buildInProcessServer(t))))
+	require.NoError(t, manager.AddClient(context.Background(), 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"))
@@ -584,7 +584,7 @@ func TestMCPGate_NoPluginsConfigured_OpStillRuns(t *testing.T) {
 	// 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))))
+	require.NoError(t, manager.AddClient(context.Background(), inProcessClientConfig("no_plugins", buildInProcessServer(t))))
 
 	clients := manager.GetClients()
 	var found bool

core/internal/mcptests/error_handling_protocol_test.go

@@ -1,6 +1,7 @@
 package mcptests
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"strings"
@@ -289,7 +290,7 @@ func TestErrorHandling_STDIO_MCPErrorResponse(t *testing.T) {
 	errorServerConfig := GetErrorTestServerConfig(bifrostRoot)
 
 	manager := setupMCPManager(t)
-	err := manager.AddClient(&errorServerConfig)
+	err := manager.AddClient(context.Background(), &errorServerConfig)
 	if err != nil {
 		t.Skipf("error-test-server not available: %v (build with: cd examples/mcps/error-test-server && go build -o bin/error-test-server)", err)
 	}
@@ -367,7 +368,7 @@ func TestErrorHandling_STDIO_TimeoutScenario(t *testing.T) {
 		ToolExecutionTimeout: schemas.Duration(2 * time.Second), // 2 second timeout
 	})
 
-	err := manager.AddClient(&errorServerConfig)
+	err := manager.AddClient(context.Background(), &errorServerConfig)
 	if err != nil {
 		t.Skipf("error-test-server not available: %v", err)
 	}
@@ -420,7 +421,7 @@ func TestErrorHandling_STDIO_MalformedJSON(t *testing.T) {
 	errorServerConfig := GetErrorTestServerConfig(bifrostRoot)
 
 	manager := setupMCPManager(t)
-	err := manager.AddClient(&errorServerConfig)
+	err := manager.AddClient(context.Background(), &errorServerConfig)
 	if err != nil {
 		t.Skipf("error-test-server not available: %v", err)
 	}
@@ -469,7 +470,7 @@ func TestErrorHandling_STDIO_IntermittentFailures(t *testing.T) {
 	errorServerConfig := GetErrorTestServerConfig(bifrostRoot)
 
 	manager := setupMCPManager(t)
-	err := manager.AddClient(&errorServerConfig)
+	err := manager.AddClient(context.Background(), &errorServerConfig)
 	if err != nil {
 		t.Skipf("error-test-server not available: %v", err)
 	}

core/internal/mcptests/health_monitoring_test.go

@@ -1,6 +1,7 @@
 package mcptests
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -50,7 +51,7 @@ func TestHealthCheckSTDIOServerDropAndRecoverIn20Seconds(t *testing.T) {
 	t.Logf("✅ Health monitor detected server drop")
 
 	// 5. Restart STDIO process (re-add client)
-	err = manager.AddClient(&clientConfig)
+	err = manager.AddClient(context.Background(), &clientConfig)
 	require.NoError(t, err, "should re-add client to simulate server recovery")
 	t.Logf("🔄 Simulated STDIO server recovery by re-adding client")
 
@@ -181,7 +182,7 @@ func TestHealthCheckStateTransitions(t *testing.T) {
 	assert.Len(t, clients, 0, "client should be removed")
 
 	// Re-add client (simulates reconnection)
-	err = manager.AddClient(&clientConfig)
+	err = manager.AddClient(context.Background(), &clientConfig)
 	require.NoError(t, err, "should re-add client")
 
 	// Verify client is connected again
@@ -394,7 +395,7 @@ func TestHealthCheckReconnectAfterFailure(t *testing.T) {
 	time.Sleep(2 * time.Second)
 
 	// Re-add client (manual reconnection)
-	err = manager.AddClient(&clientConfig)
+	err = manager.AddClient(context.Background(), &clientConfig)
 	require.NoError(t, err, "should re-add client")
 
 	// Wait for health monitoring to stabilize

core/internal/mcptests/integration_test.go

@@ -1,6 +1,7 @@
 package mcptests
 
 import (
+	"context"
 	"fmt"
 	"testing"
 	"time"
@@ -34,7 +35,7 @@ func TestIntegration_FullChatWorkflow(t *testing.T) {
 		httpConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
 		httpConfig.ID = "http-integration-test"
 		applyTestConfigHeaders(t, &httpConfig)
-		err := manager.AddClient(&httpConfig)
+		err := manager.AddClient(context.Background(), &httpConfig)
 		if err != nil {
 			t.Logf("Could not add HTTP client: %v", err)
 		}
@@ -341,7 +342,7 @@ func TestIntegration_ReconnectDuringExecution(t *testing.T) {
 	httpConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
 	httpConfig.ID = "reconnect-test-client"
 	applyTestConfigHeaders(t, &httpConfig)
-	err := manager.AddClient(&httpConfig)
+	err := manager.AddClient(context.Background(), &httpConfig)
 	require.NoError(t, err, "should add HTTP client")
 
 	// Wait for client to connect