summercms/modules/pier/stateless_fanout_test.go

package pier_test

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"

	"git.golem15.com/golem15/summercms/modules/pact"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// listenWatch is a client transport that adds the bearer and closes ended
// when the response body of the client's subscriptions/listen request ends
// (the server closed the stream).
type listenWatch struct {
	token string
	once  sync.Once
	ended chan struct{}
}

func (l *listenWatch) RoundTrip(r *http.Request) (*http.Response, error) {
	r = r.Clone(r.Context())
	r.Header.Set("Authorization", "Bearer "+l.token)
	listen := false
	if r.Body != nil {
		raw, err := io.ReadAll(r.Body)
		if err != nil {
			return nil, err
		}
		_ = r.Body.Close()
		r.Body = io.NopCloser(bytes.NewReader(raw))
		listen = strings.Contains(string(raw), `"subscriptions/listen"`)
	}
	resp, err := http.DefaultTransport.RoundTrip(r)
	if err != nil || !listen {
		return resp, err
	}
	resp.Body = &watchedBody{ReadCloser: resp.Body, done: func() { l.once.Do(func() { close(l.ended) }) }}
	return resp, nil
}

type watchedBody struct {
	io.ReadCloser
	done func()
}

func (b *watchedBody) Read(p []byte) (int, error) {
	n, err := b.ReadCloser.Read(p)
	if err != nil {
		b.done()
	}
	return n, err
}

// listenClient connects a newest-protocol client with a list_changed
// handler, so it opens a subscriptions/listen stream, and answers the
// channel closed when that stream ends.
func listenClient(t *testing.T, ctx context.Context, url, token string) (*mcp.ClientSession, <-chan struct{}) {
	t.Helper()
	watch := &listenWatch{token: token, ended: make(chan struct{})}
	client := mcp.NewClient(&mcp.Implementation{Name: "listen-test", Version: "1"}, &mcp.ClientOptions{
		ToolListChangedHandler: func(context.Context, *mcp.ToolListChangedRequest) {},
	})
	cs, err := client.Connect(ctx, &mcp.StreamableClientTransport{Endpoint: url, HTTPClient: &http.Client{Transport: watch}}, nil)
	if err != nil {
		t.Fatalf("connect with %s: %v", token, err)
	}
	t.Cleanup(func() { _ = cs.Close() })
	if cs.ID() != "" {
		t.Fatalf("the newest client got a session %q", cs.ID())
	}
	return cs, watch.ended
}

func waitEnded(t *testing.T, ended <-chan struct{}, within time.Duration, what string) {
	t.Helper()
	select {
	case <-ended:
	case <-time.After(within):
		t.Fatalf("%s: the listen stream is still open after %s", what, within)
	}
}

// TestStatelessFanOut pins plan 06's stateless lifecycle (A-WR-05,
// A-WR-06): an instance that served only stateless requests receives a
// revocation published by the other instance through Redis and evicts its
// entry; a subscriptions/listen stream ends within the re-check interval
// once its bearer stops resolving; and an LRU eviction at capacity spares a
// server with an open listen stream (RA-WR-04), evicting the least recently
// used idle one, while a revocation still ends that stream.
func TestStatelessFanOut(t *testing.T) {
	plugin := acmePlugin{tools: []pact.MCPTool{echoTool(), showTool()}}

	t.Run("a stateless-only instance applies a revocation from the other instance", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		a, b, _ := twoInstances(t, appSpec{guards: []pact.MCPGuard{&acmeGuard{}}}, plugin)
		t.Cleanup(a.pier().StopReaper)
		t.Cleanup(b.pier().StopReaper)
		cs := sdkClient(t, ctx, b.url(), "oas_at_test", clientOpts{version: "newest"})
		listNames(t, ctx, cs)
		if n := b.pier().StatelessServers(); n != 1 {
			t.Fatalf("B stateless servers = %d, want 1", n)
		}
		fireSubjectChanged(t, a.app, "acme", 7, pact.MCPSubjectRevoked)
		deadline := time.Now().Add(10 * time.Second)
		for b.pier().StatelessServers() != 0 {
			if time.Now().After(deadline) {
				t.Fatal("B never evicted its stateless server: the revocation did not reach it")
			}
			time.Sleep(20 * time.Millisecond)
		}
	})

	t.Run("a listen stream ends once its bearer stops resolving", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		verifier := newFakeVerifier(map[string]pact.MCPGrant{"oas_at_listen": grant(12, "c7", 12, "read", "write")})
		in := startInstance(t, newApp(t, appSpec{guards: []pact.MCPGuard{&acmeGuard{}}, verifier: verifier}), plugin)
		t.Cleanup(in.pier().StopReaper)
		in.pier().SetRecheck(50 * time.Millisecond)
		logs := captureLogs(t)
		_, ended := listenClient(t, ctx, in.url(), "oas_at_listen")
		select {
		case <-ended:
			t.Fatal("the listen stream ended while the bearer resolved")
		case <-time.After(300 * time.Millisecond):
		}
		verifier.revoke("oas_at_listen")
		waitEnded(t, ended, 3*time.Second, "revoked bearer")
		if !strings.Contains(logs.String(), "step=stream_closed") {
			t.Fatalf("no stream_closed line: %s", logs.String())
		}
	})

	t.Run("an LRU eviction spares a server with a listen stream", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		const capacity = 16 // max_per_binding 1 times the LRU factor
		extra := map[string]pact.MCPGrant{"oas_at_first": grant(20, "c20", 20, "read")}
		for i := 0; i < capacity; i++ {
			extra[fmt.Sprintf("oas_at_lru_%d", i)] = grant(uint(100+i), fmt.Sprintf("lru%d", i), uint(100+i), "read")
		}
		in := startInstance(t, newApp(t, appSpec{
			mcp:      "enabled: false\nsessions:\n  max_per_binding: 1\n",
			guards:   []pact.MCPGuard{&acmeGuard{}},
			verifier: newFakeVerifier(extra),
		}), plugin)
		t.Cleanup(in.pier().StopReaper)
		_, ended := listenClient(t, ctx, in.url(), "oas_at_first")
		for i := 0; i < capacity-1; i++ {
			listNames(t, ctx, sdkClient(t, ctx, in.url(), fmt.Sprintf("oas_at_lru_%d", i), clientOpts{version: "newest"}))
		}
		select {
		case <-ended:
			t.Fatal("the listen stream ended below capacity")
		case <-time.After(200 * time.Millisecond):
		}
		if n := in.pier().StatelessServers(); n != capacity {
			t.Fatalf("stateless servers at capacity = %d, want %d", n, capacity)
		}
		// One binding past capacity: the listener's server is the least
		// recently used, but it is serving the stream, so the oldest idle
		// server goes instead.
		listNames(t, ctx, sdkClient(t, ctx, in.url(), fmt.Sprintf("oas_at_lru_%d", capacity-1), clientOpts{version: "newest"}))
		select {
		case <-ended:
			t.Fatal("the LRU eviction ended a live listen stream")
		case <-time.After(300 * time.Millisecond):
		}
		if n := in.pier().StatelessServers(); n != capacity {
			t.Fatalf("stateless servers after the eviction = %d, want %d", n, capacity)
		}
		// The spared entry is still where a revocation looks for it.
		fireSubjectChanged(t, in.app, "acme", 20, pact.MCPSubjectRevoked)
		waitEnded(t, ended, 5*time.Second, "revocation after the eviction")
	})
}