summercms/modules/pier/lock_scope_test.go

package pier_test

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

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

// blockingGuard is the acme guard whose MCPAllows blocks for user 8 while
// block is set, until release closes; entered is signalled once per block.
type blockingGuard struct {
	acmeGuard
	block   atomic.Bool
	entered chan struct{}
	release chan struct{}
}

func (g *blockingGuard) MCPAllows(ctx context.Context, s pact.MCPSubject, _ pact.MCPRequirements, _ []string) (bool, error) {
	if s.UserID == 8 && g.block.Load() {
		select {
		case g.entered <- struct{}{}:
		default:
		}
		select {
		case <-g.release:
		case <-ctx.Done():
		}
	}
	return true, nil
}

// rawStatus sends one request with a bearer, an optional session id and a
// protocol version, and answers the status (safe from any goroutine).
func rawStatus(ctx context.Context, method, url, token, sid, version, body string) (int, error) {
	var rd io.Reader
	if body != "" {
		rd = bytes.NewBufferString(body)
	}
	req, err := http.NewRequestWithContext(ctx, method, url, rd)
	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", version)
	if sid != "" {
		req.Header.Set("Mcp-Session-Id", sid)
	}
	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
}

// TestLockScope pins plan 06's lock scope (A-WR-07): while one client's
// initialize is blocked in a guard check (as a database query would block
// it), another client's call on an existing session, a DELETE and a
// stateless request each finish promptly, because no router-wide lock is
// held across guard I/O.
func TestLockScope(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	guard := &blockingGuard{entered: make(chan struct{}, 1), release: make(chan struct{})}
	in := startInstance(t, newApp(t, appSpec{
		mcp:    "enabled: false\nsessions:\n  store: memory\n",
		guards: []pact.MCPGuard{guard},
	}), acmePlugin{tools: []pact.MCPTool{echoTool(), showTool()}})
	t.Cleanup(in.pier().StopReaper)

	active := openSession(t, ctx, in, "oas_at_test")
	doomed := openSession(t, ctx, in, "oas_at_test")

	guard.block.Store(true)
	blocked := make(chan int, 1)
	go func() {
		status, _ := initStatus(ctx, in.url(), "oas_at_other", initializeBody)
		blocked <- status
	}()
	select {
	case <-guard.entered:
	case <-time.After(10 * time.Second):
		t.Fatal("the blocked initialize never reached the guard")
	}

	within := func(name string, fn func() error) {
		t.Helper()
		done := make(chan error, 1)
		go func() { done <- fn() }()
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("%s: %v", name, err)
			}
		case <-time.After(3 * time.Second):
			t.Fatalf("%s waited for the blocked initialize", name)
		}
	}
	statusIs := func(want int, method, sid, version, body string) func() error {
		return func() error {
			got, err := rawStatus(ctx, method, in.url(), "oas_at_test", sid, version, body)
			if err != nil {
				return err
			}
			if got != want {
				return &statusError{got: got, want: want}
			}
			return nil
		}
	}
	t.Run("a call on an existing session", func(t *testing.T) {
		within("tools/call", statusIs(http.StatusOK, http.MethodPost, active, "2025-11-25", callBody))
	})
	t.Run("a DELETE", func(t *testing.T) {
		within("DELETE", statusIs(http.StatusNoContent, http.MethodDelete, doomed, "2025-11-25", ""))
	})
	t.Run("a stateless request", func(t *testing.T) {
		within("stateless tools/list", func() error {
			client := mcp.NewClient(&mcp.Implementation{Name: "lock-test", Version: "1"}, nil)
			transport := &mcp.StreamableClientTransport{Endpoint: in.url(), HTTPClient: &http.Client{Transport: headerTransport{token: "oas_at_test"}}}
			cs, err := client.Connect(ctx, transport, nil)
			if err != nil {
				return err
			}
			defer cs.Close()
			_, err = cs.ListTools(ctx, nil)
			return err
		})
	})

	select {
	case status := <-blocked:
		t.Fatalf("the blocked initialize finished early (%d)", status)
	default:
	}
	close(guard.release)
	select {
	case status := <-blocked:
		if status != http.StatusOK {
			t.Fatalf("the released initialize = %d", status)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("the released initialize never finished")
	}
}

type statusError struct{ got, want int }

func (e *statusError) Error() string {
	return "status " + http.StatusText(e.got) + ", want " + http.StatusText(e.want)
}