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) != 3 || 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 TestStdioSlowCancelThenEcho(t *testing.T) {
	ctx := context.Background()
	c, err := DialStdio(ctx, fakeStdioConfig(t))
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	slowCtx, cancel := context.WithTimeout(ctx, 80*time.Millisecond)
	defer cancel()
	_, err = c.CallTool(slowCtx, "slow", nil)
	if err == nil {
		t.Fatal("slow call should time out")
	}

	echoCtx, echoCancel := context.WithTimeout(ctx, 3*time.Second)
	defer echoCancel()
	res, err := c.CallTool(echoCtx, "echo", json.RawMessage(`{"text":"after-slow"}`))
	if err != nil {
		t.Fatal(err)
	}
	if len(res.Content) != 1 || res.Content[0].Text != "echo: after-slow" {
		t.Fatalf("echo must receive its own result, got %+v", res)
	}
	if c.Readers() != 1 {
		t.Fatalf("cancel must leave a single reader, got %d", c.Readers())
	}
}

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) != 3 {
		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) != 3 {
		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)
	}
}

func TestBackgroundConnectDoesNotBlockTheCaller(t *testing.T) {
	started := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-started:
		default:
			close(started)
		}
		time.Sleep(400 * time.Millisecond)
		fakeHTTPHandler(w, r)
	}))
	defer srv.Close()
	mgr, err := NewManager([]ServerConfig{{Name: "late", URL: srv.URL}})
	if err != nil {
		t.Fatal(err)
	}
	defer mgr.Close()
	reg := tools.NewRegistry(t.TempDir())
	t0 := time.Now()
	mgr.StartBackground(context.Background(), reg)
	if time.Since(t0) > 150*time.Millisecond {
		t.Fatalf("background connect blocked the caller for %s", time.Since(t0))
	}
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if reg.Has("mcp__late__echo") {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("tools did not appear after the server became ready")
}

func TestDeferredMCPToolsUseSearchAndCall(t *testing.T) {
	mgr, err := NewManager([]ServerConfig{{Name: "fake", Command: fakeStdioConfig(t).Command, Args: fakeStdioConfig(t).Args, Env: fakeStdioConfig(t).Env, Exposure: "deferred"}})
	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)
	}
	for _, d := range reg.Decls() {
		if strings.HasPrefix(d.Name, "mcp__fake__") {
			t.Fatalf("deferred mcp tool listed: %s", d.Name)
		}
	}
	if !strings.Contains(reg.Search("echo"), "mcp__fake__echo") {
		t.Fatalf("search %q", reg.Search("echo"))
	}
	res := reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"mcp__fake__echo","arguments":{"text":"via"}}`))
	if res.IsError || !strings.Contains(res.Content, "echo: via") {
		t.Fatalf("call %+v", res)
	}
}

func TestDisabledServerIsSkipped(t *testing.T) {
	off := false
	mgr, err := NewManager([]ServerConfig{{Name: "fake", Command: "x", Enabled: &off}})
	if err != nil {
		t.Fatal(err)
	}
	reg := tools.NewRegistry(t.TempDir())
	if err := mgr.RegisterInto(context.Background(), reg); err != nil {
		t.Fatal(err)
	}
	if reg.Has("mcp__fake__echo") {
		t.Fatal("disabled server registered a tool")
	}
	if err := mgr.ApplyOverrides(`{"fake":{"enabled":true,"exposure":"direct"}}`); err != nil {
		t.Fatal(err)
	}
	if !mgr.Configs()[0].IsEnabled() || mgr.Configs()[0].Exposure != "direct" {
		t.Fatalf("override %+v", mgr.Configs()[0])
	}
}

func TestManyToolsDefaultToDeferred(t *testing.T) {
	if chooseExposure("", 9) != tools.ExposureDeferred {
		t.Fatal("many tools should default to deferred")
	}
	if chooseExposure("direct", 9) != tools.ExposureDirect {
		t.Fatal("explicit direct must win")
	}
}

type askClient struct{}

func (askClient) Initialize(context.Context) (initializeResult, error) {
	return initializeResult{}, nil
}
func (askClient) ListTools(context.Context) ([]toolInfo, error) { return nil, nil }
func (askClient) Close() error                                  { return nil }
func (askClient) CallTool(context.Context, string, json.RawMessage) (callToolResult, error) {
	return callToolResult{Content: []contentBlock{{Type: "input_required", Text: "选一个颜色"}}}, nil
}

func TestInputRequiredUsesAsk(t *testing.T) {
	tool := &bridgeTool{
		client: askClient{}, declName: "mcp__x__ask", toolName: "ask",
		schema: json.RawMessage(`{"type":"object","properties":{}}`),
	}
	ctx := WithAsk(context.Background(), func(context.Context, string) (string, error) { return "蓝色", nil })
	res, err := tool.Execute(ctx, json.RawMessage(`{}`))
	if err != nil || res.IsError || res.Content != "蓝色" {
		t.Fatalf("ask result %+v %v", res, err)
	}
}