Agent-Native/runtime/mcp/mcp_test.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"agent-native/runtime/tools"
)

// TestMain：环境变量置位时，本测试二进制化身为 stdio MCP server（真实子进程测试）。
func TestMain(m *testing.M) {
	if os.Getenv(FakeServerEnv) == "1" {
		RunFakeServer()
		os.Exit(0)
	}
	os.Exit(m.Run())
}

func fakeStdioConfig(t *testing.T) ServerConfig {
	t.Helper()
	exe, err := os.Executable()
	if err != nil {
		t.Fatal(err)
	}
	return ServerConfig{
		Name:    "fake",
		Command: exe,
		Env:     append(os.Environ(), FakeServerEnv+"=1"),
	}
}

func TestStdioClientHandshakeListAndCall(t *testing.T) {
	ctx := context.Background()
	c, err := DialStdio(ctx, fakeStdioConfig(t))
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	list, err := c.ListTools(ctx)
	if err != nil || len(list) != 2 || list[0].Name != "echo" {
		t.Fatalf("list failed: %v %+v", err, list)
	}
	res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"hi"}`))
	if err != nil || len(res.Content) != 1 || res.Content[0].Text != "echo: hi" {
		t.Fatalf("call failed: %v %+v", err, res)
	}
	_, err = c.CallTool(ctx, "no_such_tool", nil)
	if err == nil || !strings.Contains(err.Error(), "unknown tool") {
		t.Fatalf("unknown-tool call must fail via rpc error, got %v", err)
	}
}

func TestStdioClientServerRejectsBadSchema(t *testing.T) {
	ctx := context.Background()
	c, err := DialStdio(ctx, fakeStdioConfig(t))
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	// server 端 tool 报 isError：作为 callToolResult 返回而非 rpc 错误。
	res, err := c.CallTool(ctx, "boom", nil)
	if err != nil {
		t.Fatal(err)
	}
	if !res.IsError || res.Content[0].Text != "boom failed" {
		t.Fatalf("want isError result, got %+v", res)
	}
}

func fakeHTTPHandler(w http.ResponseWriter, r *http.Request) {
	var req rpcRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		w.WriteHeader(400)
		return
	}
	resp := fakeServerHandle(req)
	if sid := r.Header.Get("Mcp-Session-Id"); sid != "" {
		w.Header().Set("Mcp-Session-Id", sid)
	} else {
		w.Header().Set("Mcp-Session-Id", "sess-1")
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(resp)
}

func TestHTTPClientJSONMode(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(fakeHTTPHandler))
	defer srv.Close()
	ctx := context.Background()
	c, err := DialHTTP(ctx, ServerConfig{Name: "fakehttp", URL: srv.URL})
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	if c.sessionID != "sess-1" {
		t.Fatalf("session id must be captured, got %q", c.sessionID)
	}
	list, err := c.ListTools(ctx)
	if err != nil || len(list) != 2 {
		t.Fatalf("list failed: %v %+v", err, list)
	}
	res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"net"}`))
	if err != nil || res.Content[0].Text != "echo: net" {
		t.Fatalf("call failed: %v %+v", err, res)
	}
}

func TestHTTPClientSSEMode(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req rpcRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		resp := fakeServerHandle(req)
		w.Header().Set("Content-Type", "text/event-stream")
		body, _ := json.Marshal(resp)
		fmt.Fprintf(w, "event: message\ndata: %s\n\n", body)
	}))
	defer srv.Close()
	ctx := context.Background()
	c, err := DialHTTP(ctx, ServerConfig{Name: "fakesse", URL: srv.URL})
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	list, err := c.ListTools(ctx)
	if err != nil || len(list) != 2 {
		t.Fatalf("sse list failed: %v %+v", err, list)
	}
}

func TestManagerRegisterIntoAndCallThroughRegistry(t *testing.T) {
	mgr, err := NewManager([]ServerConfig{fakeStdioConfig(t)})
	if err != nil {
		t.Fatal(err)
	}
	defer mgr.Close()
	reg := tools.NewRegistry(t.TempDir())
	if err := mgr.RegisterInto(context.Background(), reg); err != nil {
		t.Fatal(err)
	}
	var names []string
	for _, d := range reg.Decls() {
		names = append(names, d.Name)
	}
	joined := strings.Join(names, ",")
	if !strings.Contains(joined, "mcp__fake__echo") || !strings.Contains(joined, "mcp__fake__boom") {
		t.Fatalf("mcp tools not registered: %v", names)
	}
	res := reg.Execute(context.Background(), "mcp__fake__echo", json.RawMessage(`{"text":"bridge"}`))
	if res.IsError || res.Content != "echo: bridge\n" {
		t.Fatalf("bridge call failed: %+v", res)
	}
	res = reg.Execute(context.Background(), "mcp__fake__boom", nil)
	if !res.IsError || !strings.Contains(res.Content, "boom failed") {
		t.Fatalf("bridge error propagation failed: %+v", res)
	}
}

func TestManagerNamespaceWildcardRestrict(t *testing.T) {
	mgr, _ := NewManager([]ServerConfig{fakeStdioConfig(t)})
	defer mgr.Close()
	reg := tools.WorkspaceTools(t.TempDir())
	if err := mgr.RegisterInto(context.Background(), reg); err != nil {
		t.Fatal(err)
	}
	if err := reg.Restrict([]string{"mcp__fake__*"}); err != nil {
		t.Fatal(err)
	}
	for _, d := range reg.Decls() {
		if !strings.HasPrefix(d.Name, "mcp__fake__") {
			t.Fatalf("wildcard restrict must keep only mcp namespace, got %s", d.Name)
		}
	}
	// 通配未命中任何工具 → fail-closed。
	err := reg.Restrict([]string{"mcp__nope__*"})
	if err == nil || !strings.Contains(err.Error(), "mcp__nope__*") {
		t.Fatalf("unknown namespace wildcard must fail-closed, got %v", err)
	}
}

func TestManagerFailClosedOnUnreachableServer(t *testing.T) {
	mgr, err := NewManager([]ServerConfig{{Name: "dead", Command: "/nonexistent/mcp-server-bin"}})
	if err != nil {
		t.Fatal(err)
	}
	defer mgr.Close()
	reg := tools.NewRegistry(t.TempDir())
	err = mgr.RegisterInto(context.Background(), reg)
	if err == nil || !strings.Contains(err.Error(), "unreachable") {
		t.Fatalf("unreachable server must fail-closed, got %v", err)
	}
}

func TestManagerConfigValidation(t *testing.T) {
	if _, err := NewManager([]ServerConfig{{Name: "bad name", Command: "x"}}); err == nil {
		t.Fatal("invalid name must be rejected")
	}
	if _, err := NewManager([]ServerConfig{{Name: "both", Command: "x", URL: "http://y"}}); err == nil {
		t.Fatal("dual transport must be rejected")
	}
	if _, err := NewManager([]ServerConfig{{Name: "none"}}); err == nil {
		t.Fatal("missing transport must be rejected")
	}
	_, err := NewManager([]ServerConfig{{Name: "a", Command: "x"}, {Name: "a", URL: "http://y"}})
	if err == nil || !strings.Contains(err.Error(), "twice") {
		t.Fatal("duplicate names must be rejected")
	}
}

func TestManagerRegisterTimeoutBounded(t *testing.T) {
	// server 挂起不响应：initialize 必须在超时内 fail-closed，而不是永久阻塞。
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-r.Context().Done():
		case <-time.After(10 * time.Second):
		}
	}))
	defer srv.Close()
	mgr, err := NewManager([]ServerConfig{{Name: "slow", URL: srv.URL, Timeout: 300 * time.Millisecond}})
	if err != nil {
		t.Fatal(err)
	}
	defer mgr.Close()
	reg := tools.NewRegistry(t.TempDir())
	start := time.Now()
	err = mgr.RegisterInto(context.Background(), reg)
	elapsed := time.Since(start)
	if err == nil || !strings.Contains(err.Error(), "unreachable") {
		t.Fatalf("hung server must fail-closed, got %v", err)
	}
	if elapsed > 5*time.Second {
		t.Fatalf("timeout must bound the dial, took %s", elapsed)
	}
}