first commit

core/internal/mcptests/connection_test.go

@@ -0,0 +1,424 @@
+package mcptests
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/maximhq/bifrost/core/schemas"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// =============================================================================
+// HTTP CONNECTION TESTS
+// =============================================================================
+
+func TestHTTPConnection(t *testing.T) {
+	t.Parallel()
+
+	config := GetTestConfig(t)
+	if config.HTTPServerURL == "" {
+		t.Skip("MCP_HTTP_URL not set")
+	}
+
+	// Create client config
+	clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
+	// Apply headers from environment if set
+	if len(config.HTTPHeaders) > 0 {
+		clientConfig.Headers = config.HTTPHeaders
+	}
+
+	// Setup MCP manager with HTTP client
+	manager := setupMCPManager(t, clientConfig)
+
+	// Verify client was added
+	clients := manager.GetClients()
+	require.Len(t, clients, 1, "should have one client")
+	assert.Equal(t, schemas.MCPConnectionTypeHTTP, clients[0].ConnectionInfo.Type)
+	assert.Equal(t, schemas.MCPConnectionStateConnected, clients[0].State)
+}
+
+func TestHTTPConnectionInvalidURL(t *testing.T) {
+	t.Parallel()
+
+	// Create client config with invalid URL
+	invalidURL := "http://invalid-url-that-does-not-exist:9999"
+	clientConfig := GetSampleHTTPClientConfig(invalidURL)
+
+	// This should fail or have client in disconnected state
+	manager := setupMCPManager(t, clientConfig)
+	clients := manager.GetClients()
+
+	if len(clients) > 0 {
+		// If client was added, it should eventually be disconnected
+		time.Sleep(2 * time.Second)
+		clients = manager.GetClients()
+		if len(clients) > 0 {
+			assert.Equal(t, schemas.MCPConnectionStateDisconnected, clients[0].State)
+		}
+	}
+}
+
+// =============================================================================
+// SSE CONNECTION TESTS
+// =============================================================================
+
+func TestSSEConnection(t *testing.T) {
+	t.Parallel()
+
+	config := GetTestConfig(t)
+	if config.SSEServerURL == "" {
+		t.Skip("MCP_SSE_URL not set")
+	}
+
+	// Create client config
+	clientConfig := GetSampleSSEClientConfig(config.SSEServerURL)
+	// Apply headers from environment if set
+	if len(config.SSEHeaders) > 0 {
+		clientConfig.Headers = config.SSEHeaders
+	}
+
+	// Setup MCP manager with SSE client
+	manager := setupMCPManager(t, clientConfig)
+
+	// Verify client was added
+	clients := manager.GetClients()
+	require.Len(t, clients, 1, "should have one client")
+	assert.Equal(t, schemas.MCPConnectionTypeSSE, clients[0].ConnectionInfo.Type)
+	assert.Equal(t, schemas.MCPConnectionStateConnected, clients[0].State)
+}
+
+func TestSSEConnectionReconnect(t *testing.T) {
+	t.Parallel()
+
+	config := GetTestConfig(t)
+	if config.SSEServerURL == "" {
+		t.Skip("MCP_SSE_URL not set")
+	}
+
+	clientConfig := GetSampleSSEClientConfig(config.SSEServerURL)
+	// Apply headers from environment if set
+	if len(config.SSEHeaders) > 0 {
+		clientConfig.Headers = config.SSEHeaders
+	}
+	manager := setupMCPManager(t, clientConfig)
+
+	clients := manager.GetClients()
+	require.Len(t, clients, 1, "should have one client")
+
+	clientID := clients[0].ExecutionConfig.ID
+
+	// Attempt to reconnect
+	err := manager.ReconnectClient(clientID)
+	assert.NoError(t, err, "reconnect should succeed")
+
+	// Verify still connected
+	clients = manager.GetClients()
+	AssertClientState(t, clients, clientID, schemas.MCPConnectionStateConnected)
+}
+
+// =============================================================================
+// STDIO CONNECTION TESTS
+// =============================================================================
+
+func TestSTDIOConnection(t *testing.T) {
+	t.Parallel()
+
+	// Create STDIO server
+	stdioServer := NewSTDIOServerManager(t)
+	err := stdioServer.Start()
+	require.NoError(t, err, "should start STDIO server")
+	defer stdioServer.Stop()
+
+	// Wait for server to be ready
+	time.Sleep(500 * time.Millisecond)
+
+	// Note: For actual STDIO connection test, we need a compiled executable
+	// This test verifies the server manager works
+	assert.True(t, stdioServer.IsRunning(), "STDIO server should be running")
+}
+
+func TestSTDIOServerDoubleStart(t *testing.T) {
+	t.Parallel()
+
+	stdioServer := NewSTDIOServerManager(t)
+
+	// Start server
+	err := stdioServer.Start()
+	require.NoError(t, err, "first start should succeed")
+
+	// Try to start again
+	err = stdioServer.Start()
+	assert.Error(t, err, "second start should fail")
+	assert.Contains(t, err.Error(), "already running")
+}
+
+func TestSTDIOConnectionTimeout(t *testing.T) {
+	t.Parallel()
+
+	// Create client config with non-existent command
+	clientConfig := GetSampleSTDIOClientConfig("nonexistent-command", []string{})
+
+	// This should fail during connection
+	manager := setupMCPManager(t, clientConfig)
+
+	// Wait a bit for connection attempt
+	time.Sleep(2 * time.Second)
+
+	clients := manager.GetClients()
+	if len(clients) > 0 {
+		// Client should be in disconnected or error state
+		assert.NotEqual(t, schemas.MCPConnectionStateConnected, clients[0].State)
+	}
+}
+
+// =============================================================================
+// INPROCESS CONNECTION TESTS
+// =============================================================================
+
+func TestInProcessConnection(t *testing.T) {
+	t.Parallel()
+
+	// For in-process connections, we don't create a client config
+	// Instead, the internal server is created automatically when we register tools
+	manager := setupMCPManager(t)
+
+	// Register a test tool
+	toolSchema := schemas.ChatTool{
+		Type: schemas.ChatToolTypeFunction,
+		Function: &schemas.ChatToolFunction{
+			Name:        "test_inprocess_tool",
+			Description: schemas.Ptr("A test tool for in-process execution"),
+			Parameters: &schemas.ToolFunctionParameters{
+				Type: "object",
+				Properties: schemas.NewOrderedMapFromPairs(
+					schemas.KV("message", map[string]interface{}{
+						"type":        "string",
+						"description": "The message to process",
+					}),
+				),
+				Required: []string{"message"},
+			},
+		},
+	}
+	err := manager.RegisterTool(
+		"test_inprocess_tool",
+		"A test tool for in-process execution",
+		func(args any) (string, error) {
+			argsMap, ok := args.(map[string]interface{})
+			if !ok {
+				return "", assert.AnError
+			}
+			message, ok := argsMap["message"].(string)
+			if !ok {
+				return "", assert.AnError
+			}
+			result := map[string]interface{}{
+				"result": "processed: " + message,
+			}
+			resultJSON, _ := json.Marshal(result)
+			return string(resultJSON), nil
+		},
+		toolSchema,
+	)
+	require.NoError(t, err, "should register tool")
+
+	// Verify tools are available
+	ctx := createTestContext()
+	tools := manager.GetToolPerClient(ctx)
+	assert.NotEmpty(t, tools, "should have registered tool")
+}
+
+func TestInProcessToolExecution(t *testing.T) {
+	t.Parallel()
+
+	// InProcess connections don't need a client config - the internal server is created automatically
+	manager := setupMCPManager(t)
+
+	// Register a simple echo tool
+	echoToolSchema := schemas.ChatTool{
+		Type: schemas.ChatToolTypeFunction,
+		Function: &schemas.ChatToolFunction{
+			Name:        "echo_inprocess",
+			Description: schemas.Ptr("Echoes the input"),
+			Parameters: &schemas.ToolFunctionParameters{
+				Type: "object",
+				Properties: schemas.NewOrderedMapFromPairs(
+					schemas.KV("text", map[string]interface{}{
+						"type": "string",
+					}),
+				),
+			},
+		},
+	}
+	err := manager.RegisterTool(
+		"echo_inprocess",
+		"Echoes the input",
+		func(args any) (string, error) {
+			argsMap, ok := args.(map[string]interface{})
+			if !ok {
+				return "", assert.AnError
+			}
+			resultJSON, _ := json.Marshal(argsMap)
+			return string(resultJSON), nil
+		},
+		echoToolSchema,
+	)
+	require.NoError(t, err, "should register tool")
+
+	// Execute the tool
+	bifrost := setupBifrost(t)
+	bifrost.SetMCPManager(manager)
+
+	ctx := createTestContext()
+	// Create a tool call for echo_inprocess (matching the registered tool name with prefix)
+	toolCall := schemas.ChatAssistantMessageToolCall{
+		ID:   schemas.Ptr("call-1"),
+		Type: schemas.Ptr("function"),
+		Function: schemas.ChatAssistantMessageToolCallFunction{
+			Name:      schemas.Ptr("bifrostInternal-echo_inprocess"),
+			Arguments: `{"text":"test message"}`,
+		},
+	}
+
+	result, bifrostErr := bifrost.ExecuteChatMCPTool(ctx, &toolCall)
+	require.Nil(t, bifrostErr, "tool execution should succeed")
+	assert.NotNil(t, result, "should have result")
+}
+
+// =============================================================================
+// MULTIPLE CONNECTION TYPES TEST
+// =============================================================================
+
+func TestMultipleConnectionTypes(t *testing.T) {
+	t.Parallel()
+
+	config := GetTestConfig(t)
+
+	var clientConfigs []schemas.MCPClientConfig
+
+	// Add HTTP client if available
+	if config.HTTPServerURL != "" {
+		httpConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
+		httpConfig.ID = "http-client"
+		// Apply headers from environment if set
+		if len(config.HTTPHeaders) > 0 {
+			httpConfig.Headers = config.HTTPHeaders
+		}
+		clientConfigs = append(clientConfigs, httpConfig)
+	}
+
+	// Add SSE client if available
+	if config.SSEServerURL != "" {
+		sseConfig := GetSampleSSEClientConfig(config.SSEServerURL)
+		sseConfig.ID = "sse-client"
+		// Apply headers from environment if set
+		if len(config.SSEHeaders) > 0 {
+			sseConfig.Headers = config.SSEHeaders
+		}
+		clientConfigs = append(clientConfigs, sseConfig)
+	}
+
+	// Note: We don't add an InProcess client config here because InProcess connections
+	// are created automatically when tools are registered via RegisterTool()
+
+	if len(clientConfigs) == 0 {
+		t.Skip("No MCP servers configured")
+	}
+
+	// Create manager with multiple clients
+	manager := setupMCPManager(t, clientConfigs...)
+
+	// Verify all clients were added
+	clients := manager.GetClients()
+	// We expect at least the configured clients (HTTP/SSE if available)
+	assert.GreaterOrEqual(t, len(clients), len(clientConfigs), "should have all configured clients")
+
+	// Verify different connection types
+	connectionTypes := make(map[schemas.MCPConnectionType]bool)
+	for _, client := range clients {
+		connectionTypes[client.ConnectionInfo.Type] = true
+	}
+	assert.GreaterOrEqual(t, len(connectionTypes), 1, "should have at least one connection type")
+}
+
+// =============================================================================
+// CONNECTION CONFIGURATION TESTS
+// =============================================================================
+
+func TestConnectionWithHeaders(t *testing.T) {
+	t.Parallel()
+
+	config := GetTestConfig(t)
+	if config.HTTPServerURL == "" {
+		t.Skip("MCP_HTTP_URL not set")
+	}
+
+	// Create client config with custom headers
+	clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
+	clientConfig.Headers = map[string]schemas.EnvVar{
+		"Authorization":   *schemas.NewEnvVar("Bearer test-token"),
+		"X-Custom-Header": *schemas.NewEnvVar("test-value"),
+	}
+
+	manager := setupMCPManager(t, clientConfig)
+	clients := manager.GetClients()
+
+	require.Len(t, clients, 1, "should have one client")
+	assert.Equal(t, schemas.MCPConnectionStateConnected, clients[0].State)
+}
+
+func TestConnectionWithEnvironmentVariables(t *testing.T) {
+	t.Parallel()
+
+	// Create STDIO config with environment variables
+	clientConfig := GetSampleSTDIOClientConfig("echo", []string{"test"})
+	if clientConfig.StdioConfig != nil {
+		clientConfig.StdioConfig.Envs = []string{"TEST_VAR=test_value"}
+	}
+
+	// Manager creation should validate environment variables
+	manager := setupMCPManager(t, clientConfig)
+	assert.NotNil(t, manager, "should create manager")
+}
+
+func TestInvalidConnectionType(t *testing.T) {
+	t.Parallel()
+
+	// Create client config with invalid connection type
+	clientConfig := schemas.MCPClientConfig{
+		ID:             "invalid-client",
+		Name:           "Invalid Client",
+		ConnectionType: "invalid_type",
+	}
+
+	// This should fail validation
+	manager := setupMCPManager(t, clientConfig)
+
+	// Verify client was not added or is in error state
+	clients := manager.GetClients()
+	if len(clients) > 0 {
+		assert.NotEqual(t, schemas.MCPConnectionStateConnected, clients[0].State)
+	}
+}
+
+func TestConnectionWithMissingRequiredFields(t *testing.T) {
+	t.Parallel()
+
+	// HTTP connection without ConnectionString
+	clientConfig := schemas.MCPClientConfig{
+		ID:             "missing-url-client",
+		Name:           "Missing URL Client",
+		ConnectionType: schemas.MCPConnectionTypeHTTP,
+		// ConnectionString is missing
+	}
+
+	manager := setupMCPManager(t, clientConfig)
+	clients := manager.GetClients()
+
+	// Client should not be connected
+	if len(clients) > 0 {
+		assert.NotEqual(t, schemas.MCPConnectionStateConnected, clients[0].State)
+	}
+}