fn_registry/functions/infra/mcp_server_stdio_test.go

package infra

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"testing"
	"time"
)

// helper: build an MCPServerOpts wired to in/out buffers with a single echo tool.
func newTestServer(in *strings.Reader, out *bytes.Buffer) MCPServerOpts {
	echoSchema := json.RawMessage(`{"type":"object","properties":{"msg":{"type":"string"}}}`)
	return MCPServerOpts{
		Name:    "test-server",
		Version: "0.1.0",
		Tools: []MCPToolDef{
			{Name: "echo", Description: "echoes input", InputSchema: echoSchema},
		},
		Handler: func(ctx context.Context, name string, input json.RawMessage) (any, bool, error) {
			if name == "echo" {
				var args map[string]string
				_ = json.Unmarshal(input, &args)
				return map[string]string{"echoed": args["msg"]}, false, nil
			}
			return nil, true, fmt.Errorf("unknown tool: %s", name)
		},
		In:  strings.NewReader(""),
		Out: out,
	}
}

// runServer launches ServeMCP with the given lines as stdin, returns the output lines.
func runServer(t *testing.T, opts MCPServerOpts, lines []string) []map[string]any {
	t.Helper()
	payload := strings.Join(lines, "\n") + "\n"
	opts.In = strings.NewReader(payload)
	var buf bytes.Buffer
	opts.Out = &buf

	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	err := ServeMCP(ctx, opts)
	if err != nil {
		t.Fatalf("ServeMCP returned error: %v", err)
	}

	var results []map[string]any
	for _, line := range strings.Split(strings.TrimSpace(buf.String()), "\n") {
		if line == "" {
			continue
		}
		var m map[string]any
		if err2 := json.Unmarshal([]byte(line), &m); err2 != nil {
			t.Fatalf("output not valid JSON: %q — %v", line, err2)
		}
		results = append(results, m)
	}
	return results
}

func TestServeMCP_initialize(t *testing.T) {
	t.Run("initialize retorna serverInfo con Name y Version correctos", func(t *testing.T) {
		req := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`
		opts := MCPServerOpts{
			Name: "my-server", Version: "1.2.3",
			Tools:   []MCPToolDef{},
			Handler: func(_ context.Context, _ string, _ json.RawMessage) (any, bool, error) { return nil, false, nil },
		}
		results := runServer(t, opts, []string{req})
		if len(results) != 1 {
			t.Fatalf("expected 1 response, got %d", len(results))
		}
		r := results[0]
		if r["id"].(float64) != 1 {
			t.Errorf("wrong id: %v", r["id"])
		}
		result := r["result"].(map[string]any)
		info := result["serverInfo"].(map[string]any)
		if info["name"] != "my-server" {
			t.Errorf("wrong name: %v", info["name"])
		}
		if info["version"] != "1.2.3" {
			t.Errorf("wrong version: %v", info["version"])
		}
		if result["protocolVersion"] != "2024-11-05" {
			t.Errorf("wrong protocolVersion: %v", result["protocolVersion"])
		}
	})
}

func TestServeMCP_toolsList(t *testing.T) {
	t.Run("tools/list retorna las tools registradas con su schema", func(t *testing.T) {
		schema := json.RawMessage(`{"type":"object","properties":{"x":{"type":"number"}}}`)
		opts := MCPServerOpts{
			Name: "srv", Version: "0.1",
			Tools: []MCPToolDef{
				{Name: "add", Description: "adds numbers", InputSchema: schema},
			},
			Handler: func(_ context.Context, _ string, _ json.RawMessage) (any, bool, error) { return nil, false, nil },
		}
		req := `{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}`
		results := runServer(t, opts, []string{req})
		if len(results) != 1 {
			t.Fatalf("expected 1 response, got %d", len(results))
		}
		result := results[0]["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"] != "add" {
			t.Errorf("wrong tool name: %v", tool["name"])
		}
	})
}

func TestServeMCP_toolsCall_success(t *testing.T) {
	t.Run("tools/call con tool valida invoca handler y retorna content con isError false", func(t *testing.T) {
		opts := MCPServerOpts{
			Name: "srv", Version: "0.1",
			Tools: []MCPToolDef{{Name: "echo", Description: "echo", InputSchema: json.RawMessage(`{}`)}},
			Handler: func(_ context.Context, name string, input json.RawMessage) (any, bool, error) {
				var args map[string]string
				_ = json.Unmarshal(input, &args)
				return map[string]string{"echoed": args["msg"]}, false, nil
			},
		}
		req := `{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"echo","arguments":{"msg":"hello"}}}`
		results := runServer(t, opts, []string{req})
		if len(results) != 1 {
			t.Fatalf("expected 1 response, got %d", len(results))
		}
		result := results[0]["result"].(map[string]any)
		if result["isError"].(bool) {
			t.Error("expected isError=false")
		}
		content := result["content"].([]any)
		if len(content) == 0 {
			t.Fatal("expected at least 1 content block")
		}
		block := content[0].(map[string]any)
		if block["type"] != "text" {
			t.Errorf("expected type=text, got %v", block["type"])
		}
		text := block["text"].(string)
		if !strings.Contains(text, "hello") {
			t.Errorf("expected echoed hello in text, got: %s", text)
		}
	})
}

func TestServeMCP_toolsCall_handlerError(t *testing.T) {
	t.Run("tools/call cuando handler retorna error genera respuesta error -32603", func(t *testing.T) {
		opts := MCPServerOpts{
			Name: "srv", Version: "0.1",
			Tools:   []MCPToolDef{{Name: "fail", Description: "always fails", InputSchema: json.RawMessage(`{}`)}},
			Handler: func(_ context.Context, name string, _ json.RawMessage) (any, bool, error) {
				return nil, false, fmt.Errorf("internal failure")
			},
		}
		req := `{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"fail","arguments":{}}}`
		results := runServer(t, opts, []string{req})
		if len(results) != 1 {
			t.Fatalf("expected 1 response, got %d", len(results))
		}
		errField, ok := results[0]["error"].(map[string]any)
		if !ok {
			t.Fatalf("expected error field, got result: %v", results[0])
		}
		if errField["code"].(float64) != -32603 {
			t.Errorf("expected code -32603, got %v", errField["code"])
		}
	})
}

func TestServeMCP_toolsCall_isError(t *testing.T) {
	t.Run("tools/call cuando handler retorna isError=true usa result.isError=true no error JSON-RPC", func(t *testing.T) {
		opts := MCPServerOpts{
			Name: "srv", Version: "0.1",
			Tools:   []MCPToolDef{{Name: "badtool", Description: "logical error", InputSchema: json.RawMessage(`{}`)}},
			Handler: func(_ context.Context, name string, _ json.RawMessage) (any, bool, error) {
				return map[string]string{"reason": "not found"}, true, nil
			},
		}
		req := `{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"badtool","arguments":{}}}`
		results := runServer(t, opts, []string{req})
		if len(results) != 1 {
			t.Fatalf("expected 1 response, got %d", len(results))
		}
		if _, hasErr := results[0]["error"]; hasErr {
			t.Fatal("expected no JSON-RPC error field when isError=true")
		}
		result := results[0]["result"].(map[string]any)
		if !result["isError"].(bool) {
			t.Error("expected result.isError=true")
		}
	})
}

func TestServeMCP_unknownMethod(t *testing.T) {
	t.Run("metodo desconocido retorna error -32601", func(t *testing.T) {
		opts := MCPServerOpts{
			Name: "srv", Version: "0.1",
			Tools:   []MCPToolDef{},
			Handler: func(_ context.Context, _ string, _ json.RawMessage) (any, bool, error) { return nil, false, nil },
		}
		req := `{"jsonrpc":"2.0","id":6,"method":"nope/nope","params":{}}`
		results := runServer(t, opts, []string{req})
		if len(results) != 1 {
			t.Fatalf("expected 1 response, got %d", len(results))
		}
		errField, ok := results[0]["error"].(map[string]any)
		if !ok {
			t.Fatal("expected error field")
		}
		if errField["code"].(float64) != -32601 {
			t.Errorf("expected code -32601, got %v", errField["code"])
		}
	})
}

func TestServeMCP_notification(t *testing.T) {
	t.Run("notification sin id no produce respuesta en el buffer", func(t *testing.T) {
		opts := MCPServerOpts{
			Name: "srv", Version: "0.1",
			Tools:   []MCPToolDef{},
			Handler: func(_ context.Context, _ string, _ json.RawMessage) (any, bool, error) { return nil, false, nil },
		}
		// "initialized" is a notification (no id field)
		notif := `{"jsonrpc":"2.0","method":"initialized"}`
		// followed by a regular request so we know the server processed both
		req := `{"jsonrpc":"2.0","id":7,"method":"ping","params":{}}`
		results := runServer(t, opts, []string{notif, req})
		// only ping should produce a response
		if len(results) != 1 {
			t.Fatalf("expected 1 response (only for ping), got %d", len(results))
		}
		if results[0]["id"].(float64) != 7 {
			t.Errorf("expected id=7 from ping, got %v", results[0]["id"])
		}
	})
}

func TestServeMCP_invalidJSON(t *testing.T) {
	t.Run("json invalido retorna error -32700 con id null", func(t *testing.T) {
		opts := MCPServerOpts{
			Name: "srv", Version: "0.1",
			Tools:   []MCPToolDef{},
			Handler: func(_ context.Context, _ string, _ json.RawMessage) (any, bool, error) { return nil, false, nil },
		}
		req := `not json at all`
		results := runServer(t, opts, []string{req})
		if len(results) != 1 {
			t.Fatalf("expected 1 response, got %d", len(results))
		}
		errField, ok := results[0]["error"].(map[string]any)
		if !ok {
			t.Fatal("expected error field")
		}
		if errField["code"].(float64) != -32700 {
			t.Errorf("expected code -32700, got %v", errField["code"])
		}
		// id must be null (absent or nil)
		if id, hasID := results[0]["id"]; hasID && id != nil {
			t.Errorf("expected null id, got %v", id)
		}
	})
}

func TestServeMCP_ctxCancel(t *testing.T) {
	t.Run("ctx cancel detiene ServeMCP y retorna nil sin error", func(t *testing.T) {
		// Use a pipe so stdin stays open forever
		pr, pw := strings.NewReader(""), new(bytes.Buffer)
		_ = pw

		ctx, cancel := context.WithCancel(context.Background())
		opts := MCPServerOpts{
			Name: "srv", Version: "0.1",
			Tools:   []MCPToolDef{},
			Handler: func(_ context.Context, _ string, _ json.RawMessage) (any, bool, error) { return nil, false, nil },
			In:      pr,
			Out:     new(bytes.Buffer),
		}

		done := make(chan error, 1)
		go func() {
			done <- ServeMCP(ctx, opts)
		}()

		// Cancel immediately — stdin is already at EOF so it will also stop cleanly.
		cancel()

		select {
		case err := <-done:
			if err != nil {
				t.Errorf("expected nil error on ctx cancel, got: %v", err)
			}
		case <-time.After(2 * time.Second):
			t.Error("ServeMCP did not stop after ctx cancel")
		}
	})
}