summercms/modules/pier/reserve_test.go

package pier_test

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

	"git.golem15.com/golem15/summercms/modules/pact"
	"git.golem15.com/golem15/summercms/modules/pier"
)

// initStatus posts one initialize without a session and answers the status
// (safe to call from any goroutine).
func initStatus(ctx context.Context, url, token, body string) (int, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBufferString(body))
	if err != nil {
		return 0, err
	}
	req.Header.Set("Authorization", "Bearer "+token)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	req.Header.Set("Mcp-Protocol-Version", "2025-11-25")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return 0, err
	}
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<16))
	_ = resp.Body.Close()
	return resp.StatusCode, nil
}

// burst sends n concurrent initializes for token, spread over the urls, and
// counts the answers by status.
func burst(t *testing.T, ctx context.Context, n int, token string, urls ...string) map[int]int {
	t.Helper()
	var mu sync.Mutex
	got := map[int]int{}
	var wg sync.WaitGroup
	start := make(chan struct{})
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func(url string) {
			defer wg.Done()
			<-start
			status, err := initStatus(ctx, url, token, initializeBody)
			if err != nil {
				status = -1
			}
			mu.Lock()
			got[status]++
			mu.Unlock()
		}(urls[i%len(urls)])
	}
	close(start)
	wg.Wait()
	return got
}

// countingStore hides the shipped store's Reserve, Release and Update (it
// embeds only pier.Store) and counts Count and Put.
type countingStore struct {
	pier.Store
	counts atomic.Int64
	puts   atomic.Int64
}

func (s *countingStore) Count(ctx context.Context, binding string) (int, error) {
	s.counts.Add(1)
	return s.Store.Count(ctx, binding)
}

func (s *countingStore) Put(ctx context.Context, id string, rec pier.Record, ttl time.Duration) error {
	s.puts.Add(1)
	return s.Store.Put(ctx, id, rec, ttl)
}

// TestSessionReserve pins the atomic per-binding reservation of plan 06
// (A-WR-03, P-25): concurrent initializes of one binding open exactly
// mcp.sessions.max_per_binding sessions on the memory store and on a Redis
// store shared by two instances, the rest answer 429; an initialize that
// fails releases its slot; and a store without the reservation falls back
// to Count then Put.
func TestSessionReserve(t *testing.T) {
	plugin := acmePlugin{tools: []pact.MCPTool{echoTool()}}
	const cap = 3

	t.Run("memory store opens exactly max_per_binding of 50 concurrent initializes", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		in := startInstance(t, newApp(t, appSpec{
			mcp:    "enabled: false\nsessions:\n  store: memory\n  max_per_binding: 3\n",
			guards: []pact.MCPGuard{&acmeGuard{}},
		}), plugin)
		t.Cleanup(in.pier().StopReaper)
		got := burst(t, ctx, 50, "oas_at_test", in.url())
		if got[http.StatusOK] != cap || got[http.StatusTooManyRequests] != 50-cap {
			t.Fatalf("answers = %v, want %d OK and %d 429", got, cap, 50-cap)
		}
		if n := in.pier().LocalSessions(); n != cap {
			t.Fatalf("local sessions = %d, want %d", n, cap)
		}
	})

	t.Run("redis store opens exactly max_per_binding across two instances", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		a, b, _ := twoInstances(t, appSpec{
			mcp:    "enabled: false\nsessions:\n  store: redis\n  max_per_binding: 3\n",
			guards: []pact.MCPGuard{&acmeGuard{}},
		}, plugin)
		t.Cleanup(a.pier().StopReaper)
		t.Cleanup(b.pier().StopReaper)
		got := burst(t, ctx, 50, "oas_at_test", a.url(), b.url())
		if got[http.StatusOK] != cap || got[http.StatusTooManyRequests] != 50-cap {
			t.Fatalf("answers = %v, want %d OK and %d 429", got, cap, 50-cap)
		}
		if n := a.pier().LocalSessions() + b.pier().LocalSessions(); n != cap {
			t.Fatalf("local sessions on both instances = %d, want %d", n, cap)
		}
	})

	t.Run("a failed initialize releases its slot", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		in := startInstance(t, newApp(t, appSpec{
			mcp:    "enabled: false\nsessions:\n  store: memory\n  max_per_binding: 1\n",
			guards: []pact.MCPGuard{&acmeGuard{}},
		}), plugin)
		t.Cleanup(in.pier().StopReaper)
		bad := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":"not an object"}`
		// Each one reserves a slot (it gets a session id) and fails in the
		// SDK's initialize; with max_per_binding 1, a kept slot would make
		// the next one 429.
		// The SDK streams the error answer before serveNew releases the
		// slot, so each attempt waits for the release (the local copy gone).
		released := func(i int) {
			t.Helper()
			deadline := time.Now().Add(5 * time.Second)
			for in.pier().LocalSessions() != 0 {
				if time.Now().After(deadline) {
					t.Fatalf("failed initialize %d kept its local session", i)
				}
				time.Sleep(5 * time.Millisecond)
			}
		}
		for i := 0; i < 3; i++ {
			r := rpc(t, ctx, in.url(), "oas_at_test", "", bad)
			if r.status == http.StatusTooManyRequests {
				t.Fatalf("failed initialize %d hit the cap: an earlier one kept its slot", i)
			}
			if r.session == "" || !strings.Contains(r.body, `"error"`) {
				t.Fatalf("failed initialize %d = %d session %q %s, want a reserved id and an error", i, r.status, r.session, r.body)
			}
			released(i)
		}
		openSession(t, ctx, in, "oas_at_test")
		if third := rpc(t, ctx, in.url(), "oas_at_test", "", initializeBody); third.status != http.StatusTooManyRequests {
			t.Fatalf("second session over max_per_binding 1 = %d %s", third.status, third.body)
		}
	})

	t.Run("a store without the reservation falls back to Count and Put", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		in := startInstance(t, newApp(t, appSpec{
			mcp:    "enabled: false\nsessions:\n  store: memory\n  max_per_binding: 2\n",
			guards: []pact.MCPGuard{&acmeGuard{}},
		}), plugin)
		p := in.pier()
		p.StopReaper()
		store := &countingStore{}
		p.WrapStore(func(s pier.Store) pier.Store {
			store.Store = s
			return store
		})
		first := openSession(t, ctx, in, "oas_at_test")
		openSession(t, ctx, in, "oas_at_test")
		if over := rpc(t, ctx, in.url(), "oas_at_test", "", initializeBody); over.status != http.StatusTooManyRequests {
			t.Fatalf("third session = %d %s", over.status, over.body)
		}
		if store.counts.Load() < 3 || store.puts.Load() < 2 {
			t.Fatalf("Count called %d times, Put %d times: the fallback did not run", store.counts.Load(), store.puts.Load())
		}
		if call := rpc(t, ctx, in.url(), "oas_at_test", first, callBody); call.status != http.StatusOK {
			t.Fatalf("call on a fallback session = %d %s", call.status, call.body)
		}
	})
}