CyberStrikeAI/internal/mcp/connection_recovery_test.go

package mcp

import (
	"context"
	"errors"
	"fmt"
	"io"
	"testing"
	"time"

	"cyberstrike-ai/internal/config"

	"go.uber.org/zap"
)

func TestIsConnectionDeadError(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name string
		err  error
		want bool
	}{
		{"nil", nil, false},
		{"eof", io.EOF, true},
		{"wrapped eof", fmt.Errorf("connection closed: %w", io.EOF), true},
		{"client closing", errors.New(`calling "tools/list": client is closing: EOF`), true},
		{"connection reset", errors.New("read tcp: connection reset by peer"), true},
		{"canceled", context.Canceled, false},
		{"deadline", context.DeadlineExceeded, false},
		{"other", errors.New("invalid params"), false},
	}
	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			if got := isConnectionDeadError(tc.err); got != tc.want {
				t.Fatalf("isConnectionDeadError(%v) = %v, want %v", tc.err, got, tc.want)
			}
		})
	}
}

func TestLazySDKClient_MarkDisconnected(t *testing.T) {
	c := &lazySDKClient{status: "connected"}
	c.inner = &sdkClient{status: "connected"}
	c.markDisconnected()
	if c.IsConnected() {
		t.Fatal("expected disconnected after markDisconnected")
	}
	if c.GetStatus() != "disconnected" {
		t.Fatalf("expected status disconnected, got %s", c.GetStatus())
	}
}

func TestHandleConnectionDead_MarksLazyClientDisconnected(t *testing.T) {
	logger := zap.NewNop()
	m := NewExternalMCPManager(logger)

	name := "dead-mcp"
	cfg := config.ExternalMCPServerConfig{
		Type:              "http",
		URL:               "http://example.com/mcp",
		ExternalMCPEnable: true,
	}
	m.mu.Lock()
	m.configs[name] = cfg
	client := newLazySDKClient(cfg, logger)
	client.inner = &sdkClient{status: "connected"}
	client.status = "connected"
	m.clients[name] = client
	m.mu.Unlock()

	deadErr := errors.New(`connection closed: calling "tools/list": client is closing: EOF`)
	m.handleConnectionDead(name, client, deadErr)

	if client.IsConnected() {
		t.Fatal("expected disconnected after handleConnectionDead")
	}
	if m.GetError(name) == "" {
		t.Fatal("expected error message to be recorded")
	}
	counts := m.GetToolCounts()
	if counts[name] != 0 {
		t.Fatalf("expected tool count 0 after disconnect, got %d", counts[name])
	}
}

func TestReconnectBackoff(t *testing.T) {
	t.Parallel()
	if d := (&ExternalMCPManager{}).reconnectBackoff(0); d != 0 {
		t.Fatalf("attempt 0: got %v", d)
	}
	if d := (&ExternalMCPManager{}).reconnectBackoff(1); d != externalReconnectMinInterval {
		t.Fatalf("attempt 1: got %v", d)
	}
	if d := (&ExternalMCPManager{}).reconnectBackoff(10); d != externalReconnectMaxBackoff {
		t.Fatalf("attempt 10: got %v, want cap %v", d, externalReconnectMaxBackoff)
	}
}

func TestTryReconnect_RateLimited(t *testing.T) {
	logger := zap.NewNop()
	m := NewExternalMCPManager(logger)

	name := "rate-limited"
	m.reconnectMu.Lock()
	m.reconnectLastTry[name] = time.Now()
	m.reconnectAttempts[name] = 2
	m.reconnectMu.Unlock()

	m.tryReconnect(name)

	m.reconnectMu.Lock()
	attempts := m.reconnectAttempts[name]
	m.reconnectMu.Unlock()
	if attempts != 2 {
		t.Fatalf("rate limited reconnect should not increment attempts, got %d", attempts)
	}
}

func TestTryReconnect_SkipsWhenDisabled(t *testing.T) {
	logger := zap.NewNop()
	m := NewExternalMCPManager(logger)

	name := "disabled-mcp"
	m.mu.Lock()
	m.configs[name] = config.ExternalMCPServerConfig{
		Type:              "http",
		URL:               "http://example.com/mcp",
		ExternalMCPEnable: false,
	}
	m.mu.Unlock()

	m.tryReconnect(name)

	m.reconnectMu.Lock()
	attempts := m.reconnectAttempts[name]
	m.reconnectMu.Unlock()
	if attempts != 0 {
		t.Fatalf("disabled MCP should not increment reconnect attempts, got %d", attempts)
	}
}

func TestTryReconnect_SkipsWhenConnecting(t *testing.T) {
	logger := zap.NewNop()
	m := NewExternalMCPManager(logger)

	name := "connecting-mcp"
	cfg := config.ExternalMCPServerConfig{
		Type:              "http",
		URL:               "http://example.com/mcp",
		ExternalMCPEnable: true,
	}
	client := newLazySDKClient(cfg, logger)
	client.setStatus("connecting")

	m.mu.Lock()
	m.configs[name] = cfg
	m.clients[name] = client
	m.mu.Unlock()

	m.tryReconnect(name)

	m.reconnectMu.Lock()
	attempts := m.reconnectAttempts[name]
	m.reconnectMu.Unlock()
	if attempts != 0 {
		t.Fatalf("connecting MCP should not increment reconnect attempts, got %d", attempts)
	}
}

func TestStartClientAutoReconnect_SkipsWhenDisabled(t *testing.T) {
	logger := zap.NewNop()
	m := NewExternalMCPManager(logger)
	m.stopRefresh = make(chan struct{})

	name := "stopped"
	m.mu.Lock()
	m.configs[name] = config.ExternalMCPServerConfig{
		Type:              "http",
		URL:               "http://example.com/mcp",
		ExternalMCPEnable: false,
	}
	m.mu.Unlock()

	if err := m.startClient(name, true); err != nil {
		t.Fatalf("startClient: %v", err)
	}

	m.mu.RLock()
	cfg := m.configs[name]
	_, hasClient := m.clients[name]
	m.mu.RUnlock()
	if cfg.ExternalMCPEnable {
		t.Fatal("auto reconnect should not enable stopped MCP")
	}
	if hasClient {
		t.Fatal("auto reconnect should not create client when disabled")
	}
}

func TestOnClientConnected_ClearsReconnectState(t *testing.T) {
	m := &ExternalMCPManager{
		reconnectAttempts: map[string]int{"x": 3},
		reconnectLastTry:  map[string]time.Time{"x": time.Now()},
		reconnecting:      map[string]bool{"x": true},
	}
	m.onClientConnected("x")

	m.reconnectMu.Lock()
	defer m.reconnectMu.Unlock()
	if len(m.reconnectAttempts) != 0 || len(m.reconnectLastTry) != 0 || len(m.reconnecting) != 0 {
		t.Fatal("expected reconnect state cleared")
	}
}