Cyrene_v2/internal/capability/mcp/http_test.go

package mcp

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

// httpTestServer is a Streamable HTTP test server. Its GET handler keeps an SSE
// stream open until the request context is done (so it does not leak), while POST
// handling is supplied by postFn.
func httpTestServer(t *testing.T, postFn func(w http.ResponseWriter, r *http.Request)) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(http.StatusOK)
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			<-r.Context().Done()
			return
		}
		postFn(w, r)
	}))
}

// httpJSONResponse writes a single JSON-RPC response echo back, using the request
// id so the client matches it.
func httpJSONResponse(w http.ResponseWriter, r *http.Request) {
	body, _ := io.ReadAll(r.Body)
	var req rpcRequest
	_ = json.Unmarshal(body, &req)
	resp := `{"jsonrpc":"2.0","id":` + idJSON(req.ID) + `,"result":{"content":[{"type":"text","text":"direct"}]}}`
	w.Header().Set("Content-Type", "application/json")
	_, _ = w.Write([]byte(resp))
}

func TestHTTP_Call_DirectJSON(t *testing.T) {
	t.Parallel()
	srv := httpTestServer(t, httpJSONResponse)
	defer srv.Close()
	tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
	defer tr.Close()
	if err := tr.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	var res callToolResult
	if err := tr.Call(context.Background(), "tools/call", callToolParams{Name: "add"}, &res); err != nil {
		t.Fatalf("Call: %v", err)
	}
	if len(res.Content) != 1 || res.Content[0].Text != "direct" {
		t.Fatalf("Content = %+v, want [direct]", res.Content)
	}
}

func TestHTTP_Call_EventStream(t *testing.T) {
	t.Parallel()
	srv := httpTestServer(t, func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req rpcRequest
		_ = json.Unmarshal(body, &req)
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		f, _ := w.(http.Flusher)
		msg := `{"jsonrpc":"2.0","id":` + idJSON(req.ID) + `,"result":{"content":[{"type":"text","text":"stream"}]}}`
		_, _ = w.Write([]byte("event: message\ndata: " + msg + "\n\n"))
		f.Flush()
	})
	defer srv.Close()
	tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
	defer tr.Close()
	if err := tr.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	var res callToolResult
	if err := tr.Call(context.Background(), "tools/call", callToolParams{Name: "add"}, &res); err != nil {
		t.Fatalf("Call: %v", err)
	}
	if len(res.Content) != 1 || res.Content[0].Text != "stream" {
		t.Fatalf("Content = %+v, want [stream]", res.Content)
	}
}

func TestHTTP_ServerNotification(t *testing.T) {
	t.Parallel()
	notif := make(chan string, 1)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(http.StatusOK)
			f, _ := w.(http.Flusher)
			_, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\n\n"))
			f.Flush()
			<-r.Context().Done()
			return
		}
		w.WriteHeader(http.StatusAccepted)
	}))
	defer srv.Close()
	tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
	tr.setOnNotification(func(m string) { notif <- m })
	defer tr.Close()
	if err := tr.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	select {
	case m := <-notif:
		if m != notificationToolsListChanged {
			t.Fatalf("notification = %q, want %q", m, notificationToolsListChanged)
		}
	case <-time.After(time.Second):
		t.Fatal("no server notification received")
	}
}

func TestHTTP_SendsProtocolVersionHeader(t *testing.T) {
	t.Parallel()
	got := make(chan string, 1)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(http.StatusOK)
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			<-r.Context().Done()
			return
		}
		if v := r.Header.Get("Mcp-Protocol-Version"); v != "" {
			select {
			case got <- v:
			default:
			}
		}
		httpJSONResponse(w, r)
	}))
	defer srv.Close()
	tr := newHTTPTransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
	tr.setProtocolProvider(func() string { return "2025-11-25" })
	defer tr.Close()
	if err := tr.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	if err := tr.Call(context.Background(), "tools/list", map[string]any{}, &listToolsResult{}); err != nil {
		t.Fatalf("Call: %v", err)
	}
	select {
	case v := <-got:
		if v != "2025-11-25" {
			t.Fatalf("protocol header = %q, want 2025-11-25", v)
		}
	case <-time.After(time.Second):
		t.Fatal("Mcp-Protocol-Version header not sent")
	}
}

func TestHTTP_Call_NotStarted(t *testing.T) {
	t.Parallel()
	tr := newHTTPTransport(MCPConfig{Name: "s", URL: "http://example.invalid"}, defaultLogger())
	err := tr.Call(context.Background(), "tools/list", nil, nil)
	if !errors.Is(err, ErrNotConnected) {
		t.Fatalf("Call error = %v, want ErrNotConnected", err)
	}
}

func TestHTTP_Start_NoURL(t *testing.T) {
	t.Parallel()
	tr := newHTTPTransport(MCPConfig{Name: "s"}, defaultLogger())
	if err := tr.Start(context.Background()); !errors.Is(err, ErrNoTransport) {
		t.Fatalf("Start error = %v, want ErrNoTransport", err)
	}
}

func TestHTTP_Start_InvalidURL(t *testing.T) {
	t.Parallel()
	tr := newHTTPTransport(MCPConfig{Name: "s", URL: "ftp://x"}, defaultLogger())
	if err := tr.Start(context.Background()); err == nil {
		t.Fatal("Start accepted a non-http url, want error")
	}
}

func TestHTTP_Close_Idempotent(t *testing.T) {
	t.Parallel()
	tr := newHTTPTransport(MCPConfig{Name: "s", URL: "http://example.invalid"}, defaultLogger())
	if err := tr.Close(); err != nil {
		t.Fatalf("first Close: %v", err)
	}
	if err := tr.Close(); err != nil {
		t.Fatalf("second Close: %v", err)
	}
}