package mcpgateway

import (
	"context"
	"errors"
	"fmt"
	"maps"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.mongodb.org/mongo-driver/v2/mongo"

	"github.com/enterpilot/gomodel/ext"
	"github.com/enterpilot/gomodel/internal/core"
	"github.com/enterpilot/gomodel/internal/sharedstore"
	"github.com/enterpilot/gomodel/internal/storage"
	"github.com/enterpilot/gomodel/internal/storage/mongotest"
	"github.com/enterpilot/gomodel/internal/storage/sqlx/sqlxtest"
)

const listToolsBody = `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`

// runSharedStoreSuite runs body against the shared store on every storage
// backend: SQLite always, PostgreSQL and MongoDB when their test databases
// are configured.
func runSharedStoreSuite(t *testing.T, body func(t *testing.T, store ext.SharedStore)) {
	t.Helper()
	t.Run("sqlite", func(t *testing.T) {
		shared, err := storage.NewSQLite(storage.SQLiteConfig{Path: filepath.Join(t.TempDir(), "gomodel.db")})
		require.NoError(t, err)
		t.Cleanup(func() { _ = shared.Close() })
		body(t, openSharedStore(t, shared))
	})
	t.Run("postgresql", func(t *testing.T) {
		pool := sqlxtest.NewPostgresPool(t)
		if pool == nil {
			return // NewPostgresPool already skipped
		}
		body(t, openSharedStore(t, postgresStorage{pool: pool}))
	})
	mongotest.Run(t, func(t *testing.T, db *mongo.Database) {
		body(t, openSharedStore(t, mongoStorage{db: db}))
	})
}

func openSharedStore(t *testing.T, shared storage.Storage) ext.SharedStore {
	t.Helper()
	store, err := sharedstore.New(context.Background(), shared)
	require.NoError(t, err)
	t.Cleanup(func() { _ = store.Close() })
	return store
}

// postgresStorage and mongoStorage hand a test database to the shared store;
// the test helpers that opened it close it.
type postgresStorage struct{ pool *pgxpool.Pool }

func (postgresStorage) Close() error          { return nil }
func (s postgresStorage) Pool() *pgxpool.Pool { return s.pool }

type mongoStorage struct{ db *mongo.Database }

func (mongoStorage) Close() error                { return nil }
func (s mongoStorage) Database() *mongo.Database { return s.db }

// testClock runs ahead of the wall clock by a settable offset, so rechecks
// and expiry are reached without sleeping.
type testClock struct{ offset atomic.Int64 }

func (c *testClock) Now() time.Time { return time.Now().Add(time.Duration(c.offset.Load())) }

func (c *testClock) Advance(d time.Duration) { c.offset.Add(int64(d)) }

// memorySharedStore is an ext.SharedStore whose expiry follows a test clock.
type memorySharedStore struct {
	mu      sync.Mutex
	now     func() time.Time
	entries map[string]ext.SharedEntry
}

func newMemorySharedStore(now func() time.Time) *memorySharedStore {
	return &memorySharedStore{now: now, entries: map[string]ext.SharedEntry{}}
}

func (m *memorySharedStore) Put(_ context.Context, namespace, key string, value []byte, expiresAt time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if !expiresAt.IsZero() && !expiresAt.After(m.now()) {
		delete(m.entries, namespace+"/"+key)
		return nil
	}
	m.entries[namespace+"/"+key] = ext.SharedEntry{Key: key, Value: value, ExpiresAt: expiresAt}
	return nil
}

func (m *memorySharedStore) Get(_ context.Context, namespace, key string) ([]byte, bool, error) {
	entry, ok := m.entry(namespace, key)
	return entry.Value, ok, nil
}

func (m *memorySharedStore) Delete(_ context.Context, namespace, key string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.entries, namespace+"/"+key)
	return nil
}

func (m *memorySharedStore) List(_ context.Context, namespace string) ([]ext.SharedEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var entries []ext.SharedEntry
	for id, entry := range m.entries {
		if strings.HasPrefix(id, namespace+"/") && (entry.ExpiresAt.IsZero() || entry.ExpiresAt.After(m.now())) {
			entries = append(entries, entry)
		}
	}
	return entries, nil
}

func (m *memorySharedStore) entry(namespace, key string) (ext.SharedEntry, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	entry, ok := m.entries[namespace+"/"+key]
	if !ok || (!entry.ExpiresAt.IsZero() && !entry.ExpiresAt.After(m.now())) {
		return ext.SharedEntry{}, false
	}
	return entry, true
}

// replica is one gateway instance of a test cluster.
type replica struct {
	service *Service
	url     string
	clock   *testClock
}

func newReplica(t *testing.T, store ext.SharedStore, opts Options, clock *testClock) replica {
	t.Helper()
	opts.SharedStore = store
	service, url := newTestServiceWithOptions(t, opts)
	if clock == nil {
		clock = &testClock{}
	}
	service.now = clock.Now
	return replica{service: service, url: url, clock: clock}
}

func (r replica) binding(sessionID string) (sessionBinding, bool) {
	r.service.bindMu.Lock()
	defer r.service.bindMu.Unlock()
	binding, ok := r.service.bindings[sessionID]
	return binding, ok
}

func alphaServers(t *testing.T) map[string]ServerSpec {
	t.Helper()
	return map[string]ServerSpec{"alpha": testSpec("alpha", newTestUpstream(t, "alpha", addEchoTool("echo")), nil)}
}

func rawMCPDelete(t *testing.T, endpoint string, headers map[string]string) int {
	t.Helper()
	req, err := http.NewRequest(http.MethodDelete, endpoint, nil)
	require.NoError(t, err)

	for name, value := range headers {
		req.Header.Set(name, value)
	}
	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)

	defer resp.Body.Close()
	return resp.StatusCode
}

func withSession(headers map[string]string, sessionID string) map[string]string {
	merged := map[string]string{"Mcp-Session-Id": sessionID}
	maps.Copy(merged, headers)
	return merged
}

// A client opens its session on one replica, then the load balancer sends
// its requests to another: they carry on with the scope it opened with.
func TestSharedSessionContinuesOnAnotherReplica(t *testing.T) {
	runSharedStoreSuite(t, func(t *testing.T, store ext.SharedStore) {
		servers := map[string]ServerSpec{
			"alpha": testSpec("alpha", newTestUpstream(t, "alpha", addEchoTool("echo")), nil),
			"beta":  testSpec("beta", newTestUpstream(t, "beta", addEchoTool("search")), nil),
		}
		a := newReplica(t, store, Options{ConfigServers: servers}, nil)
		b := newReplica(t, store, Options{ConfigServers: servers}, nil)

		var target atomic.Pointer[Service]
		target.Store(a.service)
		balancer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			gatewayHandler(target.Load()).ServeHTTP(w, r)
		}))
		t.Cleanup(balancer.Close)

		session := connectClient(t, balancer.URL+"/mcp", map[string]string{
			ScopeHeader:          "alpha",
			"X-Test-Auth-Key-ID": "key-a",
		})
		target.Store(b.service)

		echo := NamespacedName("alpha", "echo")
		assert.Equal(t, []string{echo}, listToolNames(t, session), "the scope the session opened with holds on the other replica")
		result, err := session.CallTool(context.Background(), &mcp.CallToolParams{Name: echo, Arguments: map[string]any{"q": "hi"}})
		require.NoError(t, err)
		require.Len(t, result.Content, 1)
		assert.Equal(t, `echo:{"q":"hi"}`, result.Content[0].(*mcp.TextContent).Text)

		binding, ok := b.binding(session.ID())
		require.True(t, ok)
		assert.NotNil(t, binding.resumed, "the other replica hosts the session itself")
		assert.Equal(t, "key-a", binding.authKeyID)
	})
}

// A leaked session ID does not let another principal, or another endpoint,
// ride the session on a replica that did not open it.
func TestSharedSessionStaysBoundToItsPrincipal(t *testing.T) {
	runSharedStoreSuite(t, func(t *testing.T, store ext.SharedStore) {
		servers := alphaServers(t)
		a := newReplica(t, store, Options{ConfigServers: servers}, nil)
		b := newReplica(t, store, Options{ConfigServers: servers}, nil)
		owner := map[string]string{"X-Test-Auth-Key-ID": "key-a", core.UserPathHeader: "/team"}
		sessionID := initializeRawSession(t, a.url+"/mcp", owner)

		foreign := map[string]map[string]string{
			"auth key":  {"X-Test-Auth-Key-ID": "key-b", core.UserPathHeader: "/team"},
			"user path": {"X-Test-Auth-Key-ID": "key-a", core.UserPathHeader: "/other"},
			"no key":    {core.UserPathHeader: "/team"},
		}
		for name, headers := range foreign {
			assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(headers, sessionID)), name)
			assert.Equal(t, http.StatusNotFound, rawMCPDelete(t, b.url+"/mcp", withSession(headers, sessionID)), name)
		}
		assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, b.url+"/mcp/alpha", listToolsBody, withSession(owner, sessionID)), "another endpoint")
		_, resumed := b.binding(sessionID)
		assert.False(t, resumed, "a rejected request resumes nothing")

		assert.Equal(t, http.StatusOK, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(owner, sessionID)))
		assert.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(owner, sessionID)))
	})
}

// A DELETE on any replica ends the session on every replica: the one hosting
// it stops within a recheck interval, and none can resume it.
func TestSharedSessionDeleteEndsItOnEveryReplica(t *testing.T) {
	cases := []struct {
		name string
		// resumeOnB has replica B host the session before the DELETE.
		resumeOnB bool
		deleteOnB bool
	}{
		{name: "deleted where it was opened", resumeOnB: true},
		{name: "deleted where it was resumed", resumeOnB: true, deleteOnB: true},
		{name: "deleted where it is not hosted", deleteOnB: true},
	}
	runSharedStoreSuite(t, func(t *testing.T, store ext.SharedStore) {
		servers := alphaServers(t)
		a := newReplica(t, store, Options{ConfigServers: servers}, nil)
		b := newReplica(t, store, Options{ConfigServers: servers}, nil)
		owner := map[string]string{"X-Test-Auth-Key-ID": "key-a"}
		for _, tc := range cases {
			t.Run(tc.name, func(t *testing.T) {
				sessionID := initializeRawSession(t, a.url+"/mcp", owner)
				if tc.resumeOnB {
					require.Equal(t, http.StatusOK, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(owner, sessionID)))
				}
				deleter, other := a, b
				if tc.deleteOnB {
					deleter, other = b, a
				}
				require.Equal(t, http.StatusNoContent, rawMCPDelete(t, deleter.url+"/mcp", withSession(owner, sessionID)))
				assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, deleter.url+"/mcp", listToolsBody, withSession(owner, sessionID)))

				other.clock.Advance(sessionRecheckInterval)
				assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, other.url+"/mcp", listToolsBody, withSession(owner, sessionID)))
				_, bound := other.binding(sessionID)
				assert.False(t, bound)
			})
		}
	})
}

// Deleting a virtual server ends its sessions on every replica, before the
// others reload, and a virtual server recreated under the same name does not
// revive them.
func TestDeleteVirtualEndsSharedSessionsOnEveryReplica(t *testing.T) {
	runSharedStoreSuite(t, func(t *testing.T, store ext.SharedStore) {
		ctx := context.Background()
		servers := alphaServers(t)
		virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{
			"coding": {Name: "coding", Servers: []string{"alpha"}},
		}}
		options := func() Options {
			return Options{ConfigServers: servers, Store: &memoryStore{rows: map[string]ManagedServer{}}, VirtualStore: virtualStore}
		}
		a := newReplica(t, store, options(), nil)
		b := newReplica(t, store, options(), nil)
		sessionID := initializeRawSession(t, a.url+"/mcp/coding", nil)
		require.Equal(t, http.StatusOK, rawMCPStatus(t, b.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)))

		require.NoError(t, a.service.DeleteVirtual(ctx, "coding"))
		require.True(t, b.service.IsVirtual("coding"), "replica B has not reloaded yet")
		b.clock.Advance(sessionRecheckInterval)
		assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, b.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)))
		assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, a.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)))

		_, err := a.service.UpsertVirtual(ctx, ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}})
		require.NoError(t, err)
		c := newReplica(t, store, options(), nil)
		assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, c.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)),
			"the recreated virtual server does not revive the deleted one's session")
		fresh := initializeRawSession(t, a.url+"/mcp/coding", nil)
		assert.Equal(t, http.StatusOK, rawMCPStatus(t, c.url+"/mcp/coding", listToolsBody, withSession(nil, fresh)))
	})
}

// The shared record expires a session idle timeout after its last activity
// on any replica, and activity extends it at most once a touch interval.
func TestSharedSessionExpiresWhenIdle(t *testing.T) {
	clock := &testClock{}
	store := newMemorySharedStore(clock.Now)
	servers := alphaServers(t)
	a := newReplica(t, store, Options{ConfigServers: servers}, clock)
	b := newReplica(t, store, Options{ConfigServers: servers}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)
	expiry := func() time.Time {
		entry, ok := store.entry(sharedSessionNamespace, sessionKey(sessionID))
		require.True(t, ok)
		return entry.ExpiresAt
	}
	opened := expiry()
	assert.WithinDuration(t, clock.Now().Add(sessionIdleTimeout), opened, time.Second)

	clock.Advance(sessionTouchInterval / 2)
	require.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
	assert.Equal(t, opened, expiry(), "no write before a touch interval passes")
	clock.Advance(sessionTouchInterval / 2)
	require.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
	touched := expiry()
	assert.WithinDuration(t, clock.Now().Add(sessionIdleTimeout), touched, time.Second)

	clock.Advance(sessionIdleTimeout)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(nil, sessionID)),
		"an idle session cannot be resumed")
	_, ok := store.entry(sharedSessionNamespace, sessionKey(sessionID))
	assert.False(t, ok)
}

// The store keys sessions by a digest: the session ID itself, a bearer
// credential for the session, is never stored.
func TestSharedSessionStoresNoSessionID(t *testing.T) {
	store := newMemorySharedStore(time.Now)
	a := newReplica(t, store, Options{ConfigServers: alphaServers(t)}, nil)
	sessionID := initializeRawSession(t, a.url+"/mcp", map[string]string{"X-Test-Auth-Key-ID": "key-a"})

	entries, err := store.List(context.Background(), sharedSessionNamespace)
	require.NoError(t, err)
	require.Len(t, entries, 1)
	assert.NotContains(t, entries[0].Key, sessionID)
	assert.NotContains(t, string(entries[0].Value), sessionID)
	assert.Contains(t, string(entries[0].Value), `"auth_key_id":"key-a"`)
}

// Without a shared store each replica keeps its own sessions, as before.
func TestSessionsStayLocalWithoutSharedStore(t *testing.T) {
	servers := alphaServers(t)
	a := newReplica(t, nil, Options{ConfigServers: servers}, nil)
	b := newReplica(t, nil, Options{ConfigServers: servers}, nil)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)

	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
}

// A DELETE with an Mcp-Protocol-Version the gateway does not negotiate is
// rejected before it reaches the session, wherever the session is hosted: the
// session stays open, shared, and bound to its principal. A version from the
// per-request protocol on passes the check, as in the SDK handler.
func TestDeleteWithUnsupportedProtocolVersionKeepsTheSession(t *testing.T) {
	store := newMemorySharedStore(time.Now)
	servers := alphaServers(t)
	a := newReplica(t, store, Options{ConfigServers: servers}, nil)
	b := newReplica(t, store, Options{ConfigServers: servers}, nil)
	owner := map[string]string{"X-Test-Auth-Key-ID": "key-a"}
	foreign := map[string]string{"X-Test-Auth-Key-ID": "key-b"}
	withVersion := func(headers map[string]string, sessionID, version string) map[string]string {
		merged := withSession(headers, sessionID)
		merged["Mcp-Protocol-Version"] = version
		return merged
	}

	for name, resumed := range map[string]bool{"sdk-hosted": false, "resumed": true} {
		t.Run(name, func(t *testing.T) {
			sessionID := initializeRawSession(t, a.url+"/mcp", owner)
			host := a
			if resumed {
				host = b
				require.Equal(t, http.StatusOK, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(owner, sessionID)))
			}

			assert.Equal(t, http.StatusBadRequest, rawMCPDelete(t, host.url+"/mcp", withVersion(owner, sessionID, "2024-01-01")))
			_, bound := host.binding(sessionID)
			assert.True(t, bound, "the binding survives a rejected DELETE")
			_, stored := store.entry(sharedSessionNamespace, sessionKey(sessionID))
			assert.True(t, stored, "the shared record survives a rejected DELETE")
			_, closed := store.entry(sharedSessionNamespace, sessionCloseKey(sessionID))
			assert.False(t, closed, "a rejected DELETE writes no close mark")
			assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, host.url+"/mcp", listToolsBody, withSession(foreign, sessionID)))
			assert.Equal(t, http.StatusOK, rawMCPStatus(t, host.url+"/mcp", listToolsBody, withSession(owner, sessionID)))

			assert.Equal(t, http.StatusNoContent, rawMCPDelete(t, host.url+"/mcp", withVersion(owner, sessionID, "2099-01-01")))
		})
	}
}

// countingSharedStore counts the reads and writes that reach a shared store.
// While gate is set, reads wait for it to close.
type countingSharedStore struct {
	ext.SharedStore
	gate       chan struct{}
	gets, puts atomic.Int64
}

func (c *countingSharedStore) Get(ctx context.Context, namespace, key string) ([]byte, bool, error) {
	c.gets.Add(1)
	if c.gate != nil {
		<-c.gate
	}
	return c.SharedStore.Get(ctx, namespace, key)
}

func (c *countingSharedStore) Put(ctx context.Context, namespace, key string, value []byte, expiresAt time.Time) error {
	c.puts.Add(1)
	return c.SharedStore.Put(ctx, namespace, key, value, expiresAt)
}

// Concurrent requests for one session share its due recheck and touch: one
// reads the close mark and one extends the expiry, not each of them.
func TestConcurrentRequestsShareOneRecheck(t *testing.T) {
	const requests = 8
	clock := &testClock{}
	store := &countingSharedStore{SharedStore: newMemorySharedStore(clock.Now)}
	a := newReplica(t, store, Options{ConfigServers: alphaServers(t)}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)
	clock.Advance(sessionTouchInterval)
	store.gets.Store(0)
	store.puts.Store(0)
	// The store holds the first recheck open; the requests that share it
	// finish meanwhile.
	store.gate = make(chan struct{})

	finished := make(chan struct{}, requests)
	var wg sync.WaitGroup
	for id := range requests {
		// Distinct JSON-RPC IDs: the session rejects a duplicate in flight.
		body := fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"tools/list"}`, 10+id)
		wg.Go(func() {
			assert.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", body, withSession(nil, sessionID)))
			finished <- struct{}{}
		})
	}
	timeout := time.After(5 * time.Second)
	early := 0
wait:
	for range requests - 1 {
		select {
		case <-finished:
			early++
		case <-timeout:
			break wait
		}
	}
	close(store.gate)
	wg.Wait()
	assert.Equal(t, requests-1, early, "requests sharing the recheck finish while it is in flight")
	assert.Equal(t, int64(1), store.gets.Load(), "close-mark reads")
	assert.Equal(t, int64(1), store.puts.Load(), "expiry extensions")
}

// slowServers is an upstream whose "slow" tool runs until release is called;
// defer it, as deferred calls run before cleanup closes the servers. It
// signals entered as each call starts.
func slowServers(t *testing.T) (servers map[string]ServerSpec, entered chan struct{}, release func()) {
	t.Helper()
	entered = make(chan struct{}, 1)
	released := make(chan struct{})
	release = sync.OnceFunc(func() { close(released) })
	upstream := newTestUpstream(t, "alpha", func(server *mcp.Server) {
		server.AddTool(&mcp.Tool{Name: "slow", InputSchema: map[string]any{"type": "object"}},
			func(context.Context, *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
				entered <- struct{}{}
				<-released
				return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "done"}}}, nil
			})
	})
	return map[string]ServerSpec{"alpha": testSpec("alpha", upstream, nil)}, entered, release
}

// startSlowCall starts a tools/call of the slow tool and returns its status.
func startSlowCall(t *testing.T, endpoint, sessionID string, entered chan struct{}) chan int {
	t.Helper()
	status := make(chan int, 1)
	body := `{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"` + NamespacedName("alpha", "slow") + `","arguments":{}}}`
	req, err := http.NewRequest(http.MethodPost, endpoint, strings.NewReader(body))
	require.NoError(t, err)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	req.Header.Set("Mcp-Session-Id", sessionID)
	go func() {
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			status <- 0
			return
		}
		_ = resp.Body.Close()
		status <- resp.StatusCode
	}()
	select {
	case <-entered:
	case code := <-status:
		require.FailNowf(t, "slow call returned before the tool started", "status %d", code)
	case <-time.After(5 * time.Second):
		require.FailNow(t, "timed out waiting for the slow tool to start")
	}
	return status
}

// A call that outlasts the idle timeout keeps its session's shared record:
// the replica running it extends the record while it runs, so the session
// still continues on another replica.
func TestLongCallKeepsTheSessionShared(t *testing.T) {
	clock := &testClock{}
	store := newMemorySharedStore(clock.Now)
	servers, entered, release := slowServers(t)
	defer release()
	a := newReplica(t, store, Options{ConfigServers: servers}, clock)
	a.service.postTouchInterval = 10 * time.Millisecond
	b := newReplica(t, store, Options{ConfigServers: servers}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)
	expiry := func() time.Time {
		entry, _ := store.entry(sharedSessionNamespace, sessionKey(sessionID))
		return entry.ExpiresAt
	}

	call := startSlowCall(t, a.url+"/mcp", sessionID, entered)
	clock.Advance(sessionIdleTimeout * 2 / 3)
	require.Eventually(t, func() bool { return expiry().After(clock.Now().Add(sessionIdleTimeout - time.Minute)) },
		5*time.Second, 5*time.Millisecond, "the call in flight extends the record")
	clock.Advance(sessionIdleTimeout * 2 / 3)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(nil, sessionID)),
		"the session continues on another replica during the call")

	release()
	assert.Equal(t, http.StatusOK, <-call)
}

// The idle timeout restarts when a call ends, and so does the shared expiry:
// a session stays resumable for a full idle timeout after its last call.
func TestCallEndExtendsTheSharedSession(t *testing.T) {
	clock := &testClock{}
	store := newMemorySharedStore(clock.Now)
	servers, entered, release := slowServers(t)
	defer release()
	a := newReplica(t, store, Options{ConfigServers: servers}, clock)
	b := newReplica(t, store, Options{ConfigServers: servers}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)

	call := startSlowCall(t, a.url+"/mcp", sessionID, entered)
	clock.Advance(sessionIdleTimeout / 2)
	release()
	require.Equal(t, http.StatusOK, <-call)
	// The response reaches the client before the handler returns and ends
	// the POST, which extends the record.
	require.Eventually(t, func() bool {
		entry, ok := store.entry(sharedSessionNamespace, sessionKey(sessionID))
		return ok && entry.ExpiresAt.After(clock.Now().Add(sessionIdleTimeout-time.Minute))
	}, 5*time.Second, 5*time.Millisecond, "the end of the call extends the record")

	clock.Advance(sessionIdleTimeout * 3 / 4)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
}

// The binding sweep keeps a session with a POST in flight, whose idle timeout
// is paused, and evicts it once it has been idle a timeout after the POST.
func TestSweepKeepsSessionsWithACallInFlight(t *testing.T) {
	clock := &testClock{}
	servers, entered, release := slowServers(t)
	defer release()
	a := newReplica(t, nil, Options{ConfigServers: servers}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)

	call := startSlowCall(t, a.url+"/mcp", sessionID, entered)
	clock.Advance(sessionIdleTimeout + 2*time.Minute)
	a.service.evictIdleBindings()
	_, bound := a.binding(sessionID)
	assert.True(t, bound, "a session with a call in flight is not idle")

	release()
	require.Equal(t, http.StatusOK, <-call)
	// The response reaches the client before the handler returns and ends
	// the POST.
	require.Eventually(t, func() bool {
		binding, ok := a.binding(sessionID)
		return ok && binding.posts == 0
	}, 5*time.Second, 5*time.Millisecond)
	a.service.evictIdleBindings()
	_, bound = a.binding(sessionID)
	assert.True(t, bound, "the idle timeout counts from the end of the call")
	clock.Advance(sessionIdleTimeout + 2*time.Minute)
	a.service.evictIdleBindings()
	_, bound = a.binding(sessionID)
	assert.False(t, bound)
}

// failingListStore is a server store whose List fails while fail is set, so
// a Reload fails.
type failingListStore struct {
	*memoryStore
	fail atomic.Bool
}

func (s *failingListStore) List(ctx context.Context) ([]ManagedServer, error) {
	if s.fail.Load() {
		return nil, errors.New("store unavailable")
	}
	return s.memoryStore.List(ctx)
}

// A virtual server deleted from the store ends its sessions on the other
// replicas even when the deleting replica fails to reload afterwards.
func TestDeleteVirtualClosesSharedSessionsWhenReloadFails(t *testing.T) {
	store := newMemorySharedStore(time.Now)
	servers := alphaServers(t)
	virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{
		"coding": {Name: "coding", Servers: []string{"alpha"}, Generation: "first"},
	}}
	serverStore := &failingListStore{memoryStore: &memoryStore{rows: map[string]ManagedServer{}}}
	a := newReplica(t, store, Options{ConfigServers: servers, Store: serverStore, VirtualStore: virtualStore}, nil)
	b := newReplica(t, store, Options{ConfigServers: servers, Store: &memoryStore{rows: map[string]ManagedServer{}}, VirtualStore: virtualStore}, nil)
	sessionID := initializeRawSession(t, b.url+"/mcp/coding", nil)

	serverStore.fail.Store(true)
	require.Error(t, a.service.DeleteVirtual(context.Background(), "coding"))
	_, marked := store.entry(sharedSessionNamespace, endpointCloseKey(virtualBindingPrefix+"coding", "first"))
	assert.True(t, marked, "the endpoint close mark is written despite the failed reload")
	b.clock.Advance(sessionRecheckInterval)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, b.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)))
}

// A virtual server recreated under the same name does not revive the deleted
// one's sessions, even when the replica that opened them runs ahead of the
// one that deleted it: its creation time, read from the store, tells the two
// apart without comparing replica clocks.
func TestRecreatedVirtualIgnoresReplicaClockSkew(t *testing.T) {
	store := newMemorySharedStore(time.Now)
	ctx := context.Background()
	servers := alphaServers(t)
	created := time.Now().Add(-time.Hour).UTC()
	virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{
		"coding": {Name: "coding", Servers: []string{"alpha"}, CreatedAt: created},
	}}
	options := func() Options {
		return Options{ConfigServers: servers, Store: &memoryStore{rows: map[string]ManagedServer{}}, VirtualStore: virtualStore}
	}
	ahead := &testClock{}
	ahead.Advance(time.Minute)
	a := newReplica(t, store, options(), ahead)
	b := newReplica(t, store, options(), nil)
	sessionID := initializeRawSession(t, a.url+"/mcp/coding", nil)

	require.NoError(t, b.service.DeleteVirtual(ctx, "coding"))
	_, err := b.service.UpsertVirtual(ctx, ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}, CreatedAt: created.Add(time.Second)})
	require.NoError(t, err)
	c := newReplica(t, store, options(), nil)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, c.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)),
		"no replica resumes the deleted virtual server's session")
	fresh := initializeRawSession(t, b.url+"/mcp/coding", nil)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, c.url+"/mcp/coding", listToolsBody, withSession(nil, fresh)),
		"a session on the recreated virtual server resumes")

	require.NoError(t, a.service.Reload(ctx))
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, a.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)),
		"the replica hosting it ends it once it serves the recreated one")
	_, bound := a.binding(sessionID)
	assert.False(t, bound)
}

// A call still running when another replica deletes its session does not keep
// the deleted session's record alive: the replica running it reads the close
// mark before extending the record, and ends the session instead, so the
// session cannot resume once the mark expires.
func TestLongCallDoesNotReviveADeletedSession(t *testing.T) {
	clock := &testClock{}
	store := newMemorySharedStore(clock.Now)
	servers, entered, release := slowServers(t)
	defer release()
	a := newReplica(t, store, Options{ConfigServers: servers}, clock)
	a.service.postTouchInterval = 10 * time.Millisecond
	b := newReplica(t, store, Options{ConfigServers: servers}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)

	hosted, ok := a.binding(sessionID)
	require.True(t, ok)
	startSlowCall(t, a.url+"/mcp", sessionID, entered)
	require.Equal(t, http.StatusNoContent, rawMCPDelete(t, b.url+"/mcp", withSession(nil, sessionID)))
	clock.Advance(sessionIdleTimeout / 2)
	require.Eventually(t, func() bool {
		_, bound := a.binding(sessionID)
		return !bound
	}, 5*time.Second, 5*time.Millisecond, "the replica running the call ends the deleted session")
	_, stored := store.entry(sharedSessionNamespace, sessionKey(sessionID))
	assert.False(t, stored, "the deleted record is not written back")
	// Closing waits for the call to drain.
	release()
	assert.Eventually(t, func() bool { return !hasSessions(hosted.server) }, 5*time.Second, 5*time.Millisecond,
		"the hosted session is closed, not only unbound")

	clock.Advance(sessionCloseMarkLifetime)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, b.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
}

// failingGetStore is a shared store whose reads fail while fail is set, and
// which counts its writes.
type failingGetStore struct {
	ext.SharedStore
	fail atomic.Bool
	puts atomic.Int64
}

func (s *failingGetStore) Get(ctx context.Context, namespace, key string) ([]byte, bool, error) {
	if s.fail.Load() {
		return nil, false, errors.New("store read timed out")
	}
	return s.SharedStore.Get(ctx, namespace, key)
}

func (s *failingGetStore) Put(ctx context.Context, namespace, key string, value []byte, expiresAt time.Time) error {
	s.puts.Add(1)
	return s.SharedStore.Put(ctx, namespace, key, value, expiresAt)
}

// When the close marks cannot be read, the replica hosting a session keeps
// serving it but does not extend its record: the session may have been
// closed elsewhere, and an extended record would outlive the close mark.
func TestUnreadableCloseMarksDoNotExtendTheSession(t *testing.T) {
	clock := &testClock{}
	store := &failingGetStore{SharedStore: newMemorySharedStore(clock.Now)}
	a := newReplica(t, store, Options{ConfigServers: alphaServers(t)}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)

	store.fail.Store(true)
	store.puts.Store(0)
	clock.Advance(sessionTouchInterval)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
	assert.Zero(t, store.puts.Load(), "no expiry extension without a close-mark read")
}

// A virtual server deleted and recreated within one second, which SQL stores
// record as the same created_at, is still a new incarnation: the deleted
// one's sessions do not resume on it, wherever the replica clocks stand.
func TestVirtualRecreatedInTheSameSecondIsANewIncarnation(t *testing.T) {
	store := newMemorySharedStore(time.Now)
	ctx := context.Background()
	servers := alphaServers(t)
	created := time.Now().Add(-time.Hour).UTC().Truncate(time.Second)
	virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{
		"coding": {Name: "coding", Servers: []string{"alpha"}, CreatedAt: created, Generation: "first"},
	}}
	options := func() Options {
		return Options{ConfigServers: servers, Store: &memoryStore{rows: map[string]ManagedServer{}}, VirtualStore: virtualStore}
	}
	ahead := &testClock{}
	ahead.Advance(time.Minute)
	a := newReplica(t, store, options(), ahead)
	b := newReplica(t, store, options(), nil)
	sessionID := initializeRawSession(t, a.url+"/mcp/coding", nil)

	require.NoError(t, b.service.DeleteVirtual(ctx, "coding"))
	_, err := b.service.UpsertVirtual(ctx, ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}, CreatedAt: created})
	require.NoError(t, err)
	c := newReplica(t, store, options(), nil)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, c.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)))
	fresh := initializeRawSession(t, b.url+"/mcp/coding", nil)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, c.url+"/mcp/coding", listToolsBody, withSession(nil, fresh)))
}

// gatedListStore is a server store whose List, and so a Reload, waits on gate
// while it is set, after signalling entered.
type gatedListStore struct {
	*memoryStore
	gate    atomic.Pointer[chan struct{}]
	entered chan struct{}
}

func (s *gatedListStore) List(ctx context.Context) ([]ManagedServer, error) {
	if gate := s.gate.Load(); gate != nil {
		s.entered <- struct{}{}
		<-*gate
	}
	return s.memoryStore.List(ctx)
}

// A delete whose reload is slow closes only the incarnation it deleted: a
// server another replica recreates under the name meanwhile keeps its
// sessions, and sessions the deleting replica opens on the deleted
// incarnation while it reloads end with it.
func TestSlowVirtualDeleteSparesTheRecreatedServer(t *testing.T) {
	store := newMemorySharedStore(time.Now)
	ctx := context.Background()
	servers := alphaServers(t)
	virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{
		"coding": {Name: "coding", Servers: []string{"alpha"}, Generation: "first"},
	}}
	plainStore := func() *memoryStore { return &memoryStore{rows: map[string]ManagedServer{}} }
	gated := &gatedListStore{memoryStore: plainStore(), entered: make(chan struct{}, 1)}
	a := newReplica(t, store, Options{ConfigServers: servers, Store: gated, VirtualStore: virtualStore}, nil)
	b := newReplica(t, store, Options{ConfigServers: servers, Store: plainStore(), VirtualStore: virtualStore}, &testClock{})
	before := initializeRawSession(t, a.url+"/mcp/coding", nil)

	gate := make(chan struct{})
	gated.gate.Store(&gate)
	deleted := make(chan error, 1)
	go func() { deleted <- a.service.DeleteVirtual(ctx, "coding") }()
	<-gated.entered
	_, err := b.service.UpsertVirtual(ctx, ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}})
	require.NoError(t, err)
	fresh := initializeRawSession(t, b.url+"/mcp/coding", nil)
	during := initializeRawSession(t, a.url+"/mcp/coding", nil)
	gated.gate.Store(nil)
	close(gate)
	require.NoError(t, <-deleted)

	b.clock.Advance(sessionRecheckInterval)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, b.url+"/mcp/coding", listToolsBody, withSession(nil, fresh)),
		"the recreated server keeps its session on the replica hosting it")
	c := newReplica(t, store, Options{ConfigServers: servers, Store: plainStore(), VirtualStore: virtualStore}, nil)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, c.url+"/mcp/coding", listToolsBody, withSession(nil, fresh)),
		"and on any other")
	for name, sessionID := range map[string]string{"opened before the delete": before, "opened while it reloaded": during} {
		assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, c.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)), name)
		_, bound := a.binding(sessionID)
		assert.False(t, bound, name)
	}
}

// staleGetVirtualStore answers GetVirtual with a row another instance has
// since replaced, as a read just before the replacement would.
type staleGetVirtualStore struct {
	*memoryVirtualStore
	stale ManagedVirtualServer
}

func (s staleGetVirtualStore) GetVirtual(context.Context, string) (*ManagedVirtualServer, error) {
	return &s.stale, nil
}

// Deleting a virtual server closes the incarnation the store actually
// removed, even when the row was replaced just before the delete.
func TestDeleteVirtualClosesTheIncarnationItRemoved(t *testing.T) {
	store := newMemorySharedStore(time.Now)
	servers := alphaServers(t)
	current := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{
		"coding": {Name: "coding", Servers: []string{"alpha"}, Generation: "second"},
	}}
	plainStore := func() *memoryStore { return &memoryStore{rows: map[string]ManagedServer{}} }
	stale := staleGetVirtualStore{memoryVirtualStore: current, stale: ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}, Generation: "first"}}
	a := newReplica(t, store, Options{ConfigServers: servers, Store: plainStore(), VirtualStore: stale}, nil)
	b := newReplica(t, store, Options{ConfigServers: servers, Store: plainStore(), VirtualStore: current}, &testClock{})
	sessionID := initializeRawSession(t, b.url+"/mcp/coding", nil)

	require.NoError(t, a.service.DeleteVirtual(context.Background(), "coding"))
	b.clock.Advance(sessionRecheckInterval)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, b.url+"/mcp/coding", listToolsBody, withSession(nil, sessionID)))
}

// A replica that cannot read the close marks keeps serving a session it
// hosts only for a bounded time, shorter than a close mark lives. A session
// deleted elsewhere during a long outage therefore ends on its host, and is
// not written back once the store is readable again after the mark expired.
func TestHostEndsASessionItCannotRecheckForLong(t *testing.T) {
	clock := &testClock{}
	shared := newMemorySharedStore(clock.Now)
	unreadable := &failingGetStore{SharedStore: shared}
	servers := alphaServers(t)
	a := newReplica(t, unreadable, Options{ConfigServers: servers}, clock)
	b := newReplica(t, shared, Options{ConfigServers: servers}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)

	unreadable.fail.Store(true)
	require.Equal(t, http.StatusNoContent, rawMCPDelete(t, b.url+"/mcp", withSession(nil, sessionID)))
	clock.Advance(sessionUnverifiedLimit / 2)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)),
		"a short outage keeps the session on its host")
	clock.Advance(sessionUnverifiedLimit / 2)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)),
		"the host ends a session it has not rechecked for the limit")

	clock.Advance(sessionCloseMarkLifetime)
	unreadable.fail.Store(false)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
	_, stored := shared.entry(sharedSessionNamespace, sessionKey(sessionID))
	assert.False(t, stored, "the deleted record is not written back after recovery")
}

// The same bound holds while a call runs: the replica running it ends the
// session once its close marks have been unreadable for the limit.
func TestLongCallEndsASessionItCannotRecheckForLong(t *testing.T) {
	clock := &testClock{}
	unreadable := &failingGetStore{SharedStore: newMemorySharedStore(clock.Now)}
	servers, entered, release := slowServers(t)
	defer release()
	a := newReplica(t, unreadable, Options{ConfigServers: servers}, clock)
	a.service.postTouchInterval = 10 * time.Millisecond
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)

	startSlowCall(t, a.url+"/mcp", sessionID, entered)
	unreadable.fail.Store(true)
	clock.Advance(sessionUnverifiedLimit)
	assert.Eventually(t, func() bool {
		_, bound := a.binding(sessionID)
		return !bound
	}, 5*time.Second, 5*time.Millisecond)
}

// The bound counts from the last successful recheck, not from when the
// session opened: an outage that starts late in a long session gets the
// whole limit.
func TestUnverifiedLimitCountsFromTheLastRecheck(t *testing.T) {
	clock := &testClock{}
	unreadable := &failingGetStore{SharedStore: newMemorySharedStore(clock.Now)}
	a := newReplica(t, unreadable, Options{ConfigServers: alphaServers(t)}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)

	clock.Advance(sessionUnverifiedLimit * 3 / 4)
	require.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
	unreadable.fail.Store(true)
	clock.Advance(sessionUnverifiedLimit * 3 / 4)
	assert.Equal(t, http.StatusOK, rawMCPStatus(t, a.url+"/mcp", listToolsBody, withSession(nil, sessionID)))
}

// A failed recheck applies the limit to the session's current verification,
// not the snapshot the request started with: a recheck that succeeded
// meanwhile keeps the session.
func TestFailedRecheckSeesAConcurrentVerification(t *testing.T) {
	clock := &testClock{}
	unreadable := &failingGetStore{SharedStore: newMemorySharedStore(clock.Now)}
	a := newReplica(t, unreadable, Options{ConfigServers: alphaServers(t)}, clock)
	sessionID := initializeRawSession(t, a.url+"/mcp", nil)
	snapshot, ok := a.binding(sessionID)
	require.True(t, ok)

	clock.Advance(sessionUnverifiedLimit)
	a.service.markVerified(sessionID, snapshot.server)
	unreadable.fail.Store(true)
	assert.True(t, a.service.recheckSharedSession(context.Background(), sessionID, snapshot))
}