server_test.go

package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
)

// testProvider implements Provider for testing
type testProvider struct {
	resources []ResourceHandler
	tools     []ToolHandler
}

func (p *testProvider) Resources() []ResourceHandler { return p.resources }
func (p *testProvider) Tools() []ToolHandler          { return p.tools }

func sendRequest(t *testing.T, server *Server, method string, params any) map[string]any {
	t.Helper()
	req := map[string]any{
		"jsonrpc": "2.0",
		"id":      1,
		"method":  method,
	}
	if params != nil {
		raw, _ := json.Marshal(params)
		req["params"] = json.RawMessage(raw)
	}
	body, _ := json.Marshal(req)

	rr := httptest.NewRecorder()
	httpReq := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body))
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("X-Workspace-ID", "test-workspace")

	server.ServeHTTP(rr, httpReq)

	if rr.Code != http.StatusOK {
		t.Fatalf("expected status 200, got %d", rr.Code)
	}

	var resp map[string]any
	if err := json.NewDecoder(rr.Body).Decode(&resp); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	return resp
}

func TestInitialize(t *testing.T) {
	server := NewServer("test-server", "1.0.0")

	resp := sendRequest(t, server, "initialize", map[string]any{
		"protocolVersion": "2024-11-05",
		"capabilities":    map[string]any{},
		"clientInfo":      map[string]any{"name": "test", "version": "1.0.0"},
	})

	result, ok := resp["result"].(map[string]any)
	if !ok {
		t.Fatalf("expected result object, got %T", resp["result"])
	}

	if result["protocolVersion"] != "2024-11-05" {
		t.Errorf("expected protocol version 2024-11-05, got %v", result["protocolVersion"])
	}

	serverInfo := result["serverInfo"].(map[string]any)
	if serverInfo["name"] != "test-server" {
		t.Errorf("expected server name test-server, got %v", serverInfo["name"])
	}
}

func TestInitializeWithTools(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	server.AddProvider(&testProvider{
		tools: []ToolHandler{
			{
				Definition: Tool{Name: "test_tool", Description: "A test tool", InputSchema: json.RawMessage(`{"type":"object"}`)},
				Call: func(ctx context.Context, workspaceID string, arguments json.RawMessage) (*CallToolResult, error) {
					return &CallToolResult{Content: []ToolContent{{Type: "text", Text: "ok"}}}, nil
				},
			},
		},
	})

	resp := sendRequest(t, server, "initialize", map[string]any{
		"protocolVersion": "2024-11-05",
		"capabilities":    map[string]any{},
		"clientInfo":      map[string]any{"name": "test", "version": "1.0.0"},
	})

	result := resp["result"].(map[string]any)
	caps := result["capabilities"].(map[string]any)

	if caps["resources"] == nil {
		t.Error("expected resources capability")
	}
	if caps["tools"] == nil {
		t.Error("expected tools capability when tools are registered")
	}
}

func TestInitializeWithoutTools(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	server.AddProvider(&testProvider{
		resources: []ResourceHandler{
			{Scheme: "test", Name: "Test", List: func(ctx context.Context, wid string) ([]Resource, error) {
				return nil, nil
			}, Read: func(ctx context.Context, wid, uri string) (*ResourceContent, error) {
				return nil, nil
			}},
		},
	})

	resp := sendRequest(t, server, "initialize", map[string]any{
		"protocolVersion": "2024-11-05",
		"capabilities":    map[string]any{},
		"clientInfo":      map[string]any{"name": "test", "version": "1.0.0"},
	})

	result := resp["result"].(map[string]any)
	caps := result["capabilities"].(map[string]any)

	if caps["tools"] != nil {
		t.Error("expected no tools capability when no tools registered")
	}
}

func TestPing(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	resp := sendRequest(t, server, "ping", nil)

	if resp["error"] != nil {
		t.Errorf("expected no error, got %v", resp["error"])
	}
}

func TestNotificationsInitialized(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	resp := sendRequest(t, server, "notifications/initialized", nil)

	if resp["error"] != nil {
		t.Errorf("expected no error, got %v", resp["error"])
	}
}

func TestMethodNotFound(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	resp := sendRequest(t, server, "unknown/method", nil)

	errObj, ok := resp["error"].(map[string]any)
	if !ok {
		t.Fatalf("expected error object")
	}
	if errObj["code"].(float64) != -32601 {
		t.Errorf("expected error code -32601, got %v", errObj["code"])
	}
}

func TestProviderRegistration(t *testing.T) {
	server := NewServer("test-server", "1.0.0")

	provider := &testProvider{
		resources: []ResourceHandler{
			{
				Scheme: "test",
				Name:   "Test Resources",
				List: func(ctx context.Context, workspaceID string) ([]Resource, error) {
					return []Resource{{URI: "test://1", Name: "Item 1", MimeType: "application/json"}}, nil
				},
				Read: func(ctx context.Context, workspaceID, uri string) (*ResourceContent, error) {
					return &ResourceContent{URI: uri, MimeType: "application/json", Text: `{"id":"1"}`}, nil
				},
			},
		},
		tools: []ToolHandler{
			{
				Definition: Tool{Name: "do_thing", Description: "Does a thing", InputSchema: json.RawMessage(`{"type":"object"}`)},
				Call: func(ctx context.Context, workspaceID string, arguments json.RawMessage) (*CallToolResult, error) {
					return &CallToolResult{Content: []ToolContent{{Type: "text", Text: "done"}}}, nil
				},
			},
		},
	}
	server.AddProvider(provider)

	// Verify resources
	resp := sendRequest(t, server, "resources/list", nil)
	result := resp["result"].(map[string]any)
	resources := result["resources"].([]any)
	if len(resources) != 1 {
		t.Errorf("expected 1 resource, got %d", len(resources))
	}

	// Verify tools
	resp = sendRequest(t, server, "tools/list", nil)
	result = resp["result"].(map[string]any)
	tools := result["tools"].([]any)
	if len(tools) != 1 {
		t.Errorf("expected 1 tool, got %d", len(tools))
	}
}

func TestMultipleProviders(t *testing.T) {
	server := NewServer("test-server", "1.0.0")

	server.AddProvider(&testProvider{
		resources: []ResourceHandler{
			{Scheme: "a", Name: "A", List: func(ctx context.Context, wid string) ([]Resource, error) {
				return []Resource{{URI: "a://1", Name: "A1"}}, nil
			}, Read: func(ctx context.Context, wid, uri string) (*ResourceContent, error) {
				return nil, nil
			}},
		},
		tools: []ToolHandler{
			{Definition: Tool{Name: "tool_a", InputSchema: json.RawMessage(`{"type":"object"}`)},
				Call: func(ctx context.Context, wid string, args json.RawMessage) (*CallToolResult, error) {
					return &CallToolResult{Content: []ToolContent{{Type: "text", Text: "a"}}}, nil
				}},
		},
	})

	server.AddProvider(&testProvider{
		resources: []ResourceHandler{
			{Scheme: "b", Name: "B", List: func(ctx context.Context, wid string) ([]Resource, error) {
				return []Resource{{URI: "b://1", Name: "B1"}}, nil
			}, Read: func(ctx context.Context, wid, uri string) (*ResourceContent, error) {
				return nil, nil
			}},
		},
		tools: []ToolHandler{
			{Definition: Tool{Name: "tool_b", InputSchema: json.RawMessage(`{"type":"object"}`)},
				Call: func(ctx context.Context, wid string, args json.RawMessage) (*CallToolResult, error) {
					return &CallToolResult{Content: []ToolContent{{Type: "text", Text: "b"}}}, nil
				}},
		},
	})

	// Should have 2 resources
	resp := sendRequest(t, server, "resources/list", nil)
	result := resp["result"].(map[string]any)
	resources := result["resources"].([]any)
	if len(resources) != 2 {
		t.Errorf("expected 2 resources from 2 providers, got %d", len(resources))
	}

	// Should have 2 tools
	resp = sendRequest(t, server, "tools/list", nil)
	result = resp["result"].(map[string]any)
	tools := result["tools"].([]any)
	if len(tools) != 2 {
		t.Errorf("expected 2 tools from 2 providers, got %d", len(tools))
	}
}

func TestToolsList(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	server.AddProvider(&testProvider{
		tools: []ToolHandler{
			{
				Definition: Tool{
					Name:        "get_item",
					Description: "Gets an item by ID",
					InputSchema: json.RawMessage(`{"type":"object","properties":{"id":{"type":"string"}},"required":["id"]}`),
				},
				Call: func(ctx context.Context, wid string, args json.RawMessage) (*CallToolResult, error) {
					return nil, nil
				},
			},
		},
	})

	resp := sendRequest(t, server, "tools/list", nil)
	result := resp["result"].(map[string]any)
	tools := result["tools"].([]any)

	if len(tools) != 1 {
		t.Fatalf("expected 1 tool, got %d", len(tools))
	}

	tool := tools[0].(map[string]any)
	if tool["name"] != "get_item" {
		t.Errorf("expected tool name get_item, got %v", tool["name"])
	}
	if tool["description"] != "Gets an item by ID" {
		t.Errorf("expected tool description, got %v", tool["description"])
	}
	if tool["inputSchema"] == nil {
		t.Error("expected inputSchema")
	}
}

func TestToolsCallSuccess(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	server.AddProvider(&testProvider{
		tools: []ToolHandler{
			{
				Definition: Tool{Name: "greet", InputSchema: json.RawMessage(`{"type":"object","properties":{"name":{"type":"string"}}}`)},
				Call: func(ctx context.Context, wid string, args json.RawMessage) (*CallToolResult, error) {
					var a struct{ Name string `json:"name"` }
					json.Unmarshal(args, &a)
					return &CallToolResult{
						Content: []ToolContent{{Type: "text", Text: "Hello, " + a.Name}},
					}, nil
				},
			},
		},
	})

	resp := sendRequest(t, server, "tools/call", map[string]any{
		"name":      "greet",
		"arguments": map[string]any{"name": "World"},
	})

	if resp["error"] != nil {
		t.Fatalf("unexpected error: %v", resp["error"])
	}

	result := resp["result"].(map[string]any)
	content := result["content"].([]any)
	if len(content) != 1 {
		t.Fatalf("expected 1 content item, got %d", len(content))
	}

	item := content[0].(map[string]any)
	if item["text"] != "Hello, World" {
		t.Errorf("expected 'Hello, World', got %v", item["text"])
	}
}

func TestToolsCallError(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	server.AddProvider(&testProvider{
		tools: []ToolHandler{
			{
				Definition: Tool{Name: "fail_tool", InputSchema: json.RawMessage(`{"type":"object"}`)},
				Call: func(ctx context.Context, wid string, args json.RawMessage) (*CallToolResult, error) {
					return nil, context.DeadlineExceeded
				},
			},
		},
	})

	resp := sendRequest(t, server, "tools/call", map[string]any{
		"name": "fail_tool",
	})

	// Tool errors are returned as CallToolResult with IsError, not as JSON-RPC errors
	if resp["error"] != nil {
		t.Fatalf("unexpected JSON-RPC error: %v", resp["error"])
	}

	result := resp["result"].(map[string]any)
	if result["isError"] != true {
		t.Error("expected isError to be true")
	}

	content := result["content"].([]any)
	item := content[0].(map[string]any)
	if item["text"] != "error: context deadline exceeded" {
		t.Errorf("expected error message, got %v", item["text"])
	}
}

func TestToolsCallUnknownTool(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	server.AddProvider(&testProvider{
		tools: []ToolHandler{
			{
				Definition: Tool{Name: "real_tool", InputSchema: json.RawMessage(`{"type":"object"}`)},
				Call: func(ctx context.Context, wid string, args json.RawMessage) (*CallToolResult, error) {
					return nil, nil
				},
			},
		},
	})

	resp := sendRequest(t, server, "tools/call", map[string]any{
		"name": "nonexistent_tool",
	})

	errObj, ok := resp["error"].(map[string]any)
	if !ok {
		t.Fatal("expected error for unknown tool")
	}
	if errObj["code"].(float64) != -32601 {
		t.Errorf("expected method not found error code, got %v", errObj["code"])
	}
}

func TestToolsCallMissingParams(t *testing.T) {
	server := NewServer("test-server", "1.0.0")
	server.AddProvider(&testProvider{
		tools: []ToolHandler{
			{
				Definition: Tool{Name: "test_tool", InputSchema: json.RawMessage(`{"type":"object"}`)},
				Call: func(ctx context.Context, wid string, args json.RawMessage) (*CallToolResult, error) {
					return nil, nil
				},
			},
		},
	})

	// Send tools/call with no params
	req := map[string]any{
		"jsonrpc": "2.0",
		"id":      1,
		"method":  "tools/call",
	}
	body, _ := json.Marshal(req)

	rr := httptest.NewRecorder()
	httpReq := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body))
	httpReq.Header.Set("Content-Type", "application/json")

	server.ServeHTTP(rr, httpReq)

	var resp map[string]any
	json.NewDecoder(rr.Body).Decode(&resp)

	errObj, ok := resp["error"].(map[string]any)
	if !ok {
		t.Fatal("expected error for missing params")
	}
	if errObj["code"].(float64) != -32602 {
		t.Errorf("expected invalid params error code, got %v", errObj["code"])
	}
}