Cyrene_v2/internal/capability/mcp/sse_test.go

package mcp

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

// postRespondServer serves a GET stream that ends immediately and answers every
// POST directly with a JSON-RPC response body (the non-streaming reply path).
func postRespondServer(t *testing.T) *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()
			}
			return
		}
		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))
	}))
}

// streamServer pushes responses to the POST caller back through the SSE stream.
func streamServer(t *testing.T) *httptest.Server {
	t.Helper()
	push := make(chan string, 8)
	var srv *httptest.Server
	srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet:
			w.Header().Set("Content-Type", "text/event-stream")
			w.Header().Set("Cache-Control", "no-cache")
			f, _ := w.(http.Flusher)
			_, _ = w.Write([]byte("event: endpoint\ndata: " + srv.URL + "\n\n"))
			f.Flush()
			for {
				select {
				case data := <-push:
					_, _ = w.Write([]byte("event: message\ndata: " + data + "\n\n"))
					f.Flush()
				case <-r.Context().Done():
					return
				}
			}
		default:
			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":"stream"}]}}`
			push <- resp
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("ok"))
		}
	}))
	return srv
}

// idJSON renders a request id as a JSON string literal when it is a string.
func idJSON(id any) string {
	if s, ok := id.(string); ok {
		b, _ := json.Marshal(s)
		return string(b)
	}
	return `null`
}

func TestSSE_Call_DirectResponse(t *testing.T) {
	t.Parallel()
	srv := postRespondServer(t)
	defer srv.Close()
	tr := newSSETransport(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", Arguments: map[string]any{}}, &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 TestSSE_Call_MessageEvent(t *testing.T) {
	t.Parallel()
	srv := streamServer(t)
	defer srv.Close()
	tr := newSSETransport(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", Arguments: map[string]any{}}, &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 TestSSE_Start_Non200(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	defer srv.Close()
	tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
	if err := tr.Start(context.Background()); err == nil {
		t.Fatal("Start succeeded on a non-200 response, want error")
	}
}

func TestSSE_Call_NotStarted(t *testing.T) {
	t.Parallel()
	tr := newSSETransport(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 TestSSE_Call_ContextTimeout(t *testing.T) {
	t.Parallel()
	// A server whose GET stream stays open and whose POST replies with no body
	// and never pushes forces the Call to wait on its context deadline.
	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
		}
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))
	defer srv.Close()

	tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
	if err := tr.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer tr.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	err := tr.Call(ctx, "tools/call", callToolParams{Name: "add"}, &callToolResult{})
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("Call error = %v, want context.DeadlineExceeded", err)
	}
}

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

func TestSSE_Close_ClosesPendingAndBody(t *testing.T) {
	t.Parallel()
	srv := postRespondServer(t)
	defer srv.Close()
	tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
	if err := tr.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	// Put a sentinel pending call in place.
	tr.mu.Lock()
	tr.pending["999"] = make(chan []byte, 1)
	tr.mu.Unlock()
	if err := tr.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
	tr.mu.Lock()
	empty := len(tr.pending) == 0
	tr.mu.Unlock()
	if !empty {
		t.Fatalf("pending not cleared after Close")
	}
	if err := tr.Close(); err != nil {
		t.Fatalf("second Close: %v", err)
	}
}