package api

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/api/auth"
	"github.com/tidbcloud/AgentTeam/desktop/internal/api/mcp"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/streamregistry"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/uisession"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire"
)

type gatewayTestMCPHandler struct{}

func (gatewayTestMCPHandler) Tools() []types.MCPToolDef { return nil }
func (gatewayTestMCPHandler) CallTool(context.Context, string, string, json.RawMessage) *types.MCPToolResult {
	return nil
}
func (gatewayTestMCPHandler) SkillDoc(types.ToolNamer) string { return "" }
func (gatewayTestMCPHandler) SkillDocSplit(types.ToolNamer) types.SkillDocParts {
	return types.SkillDocParts{}
}

type gatewayTestConnection struct {
	id           string
	serverURL    string
	userToken    string
	refreshToken string
	tokenCalls   atomic.Int32
	invalidates  atomic.Int32
	httpClient   *http.Client
	tokenHook    func(context.Context) (string, error)
	refreshHook  func(context.Context, string) (string, error)

	// gatewayClient mirrors wire.Protocol's lifecycle-scoped gateway
	// client. The first GatewayHTTPClient() call constructs and
	// caches it; subsequent calls return the same instance.
	gatewayClient   *http.Client
	gatewayClientMu sync.Mutex
}

func (c *gatewayTestConnection) ID() string                  { return c.id }
func (c *gatewayTestConnection) Enabled() bool               { return true }
func (c *gatewayTestConnection) Start(context.Context) error { return nil }
func (c *gatewayTestConnection) Stop(context.Context) error  { return nil }
func (c *gatewayTestConnection) Send(context.Context, types.Metadata, types.OutboundMessage) error {
	return nil
}
func (c *gatewayTestConnection) ServerURL() string { return c.serverURL }

// GatewayHTTPClient mirrors wire.Protocol.GatewayHTTPClient by
// delegating to the production-owned wire.BuildGatewayHTTPClient
// builder. This keeps the policy shape (Timeout=0, Jar=nil,
// CheckRedirect=ErrUseLastResponse, cloned Transport with
// ResponseHeaderTimeout) identical between production and tests;
// only the ResponseHeaderTimeout duration differs (test-scoped
// short value) so nil-Transport integration tests don't have to
// wait the full 30s production bound.
func (c *gatewayTestConnection) GatewayHTTPClient() *http.Client {
	c.gatewayClientMu.Lock()
	defer c.gatewayClientMu.Unlock()
	if c.gatewayClient != nil {
		return c.gatewayClient
	}
	c.gatewayClient = wire.BuildGatewayHTTPClient(c.httpClient, gatewayTestResponseHeaderTimeout)
	return c.gatewayClient
}

// gatewayTestResponseHeaderTimeout is the test-scoped duration for
// the gateway client's ResponseHeaderTimeout. Identical in shape to
// wire.GatewayResponseHeaderTimeout but shorter so the nil-Transport
// integration test fails fast instead of waiting the production 30s.
const gatewayTestResponseHeaderTimeout = 200 * time.Millisecond

func (c *gatewayTestConnection) UserToken(ctx context.Context) (string, error) {
	if c.tokenHook != nil {
		return c.tokenHook(ctx)
	}
	call := c.tokenCalls.Add(1)
	if call > 1 && c.refreshToken != "" {
		return c.refreshToken, nil
	}
	return c.userToken, nil
}
func (c *gatewayTestConnection) RefreshUserToken(ctx context.Context, stale string) (string, error) {
	if c.refreshHook != nil {
		return c.refreshHook(ctx, stale)
	}
	c.invalidates.Add(1)
	c.tokenCalls.Add(1)
	if c.refreshToken != "" {
		return c.refreshToken, nil
	}
	return c.userToken, nil
}
func (c *gatewayTestConnection) ServerIdentitySnapshot() profile.AgentTeamServer {
	return profile.AgentTeamServer{
		ID:        c.id,
		Enabled:   true,
		ServerURL: c.serverURL,
	}
}

type gatewayTestConnectionAccessor struct {
	mu          sync.RWMutex
	connections map[string]types.ServerConnection
}

func (a *gatewayTestConnectionAccessor) ServerConnection(id string) (types.ServerConnection, bool) {
	a.mu.RLock()
	defer a.mu.RUnlock()
	connection, ok := a.connections[id]
	return connection, ok
}
func (a *gatewayTestConnectionAccessor) ServerConnections() []types.ServerConnection {
	a.mu.RLock()
	defer a.mu.RUnlock()
	result := make([]types.ServerConnection, 0, len(a.connections))
	for _, connection := range a.connections {
		result = append(result, connection)
	}
	return result
}
func (*gatewayTestConnectionAccessor) Enable(string) error  { return nil }
func (*gatewayTestConnectionAccessor) Disable(string) error { return nil }

type blockingGatewayReconciler struct {
	*gatewayTestConnectionAccessor
	entered       chan struct{}
	release       chan struct{}
	newConnection types.ServerConnection
}

func (r *blockingGatewayReconciler) Reconcile(
	_ context.Context,
	_ *profile.AgentTeamServer,
	current *profile.AgentTeamServer,
) error {
	close(r.entered)
	<-r.release
	r.mu.Lock()
	r.connections[current.ID] = r.newConnection
	r.mu.Unlock()
	return nil
}

type failingGatewayReconciler struct {
	*gatewayTestConnectionAccessor
	calls int
}

func (r *failingGatewayReconciler) Reconcile(
	_ context.Context,
	_ *profile.AgentTeamServer,
	_ *profile.AgentTeamServer,
) error {
	r.calls++
	if r.calls == 1 {
		return errors.New("reconcile failed")
	}
	return nil
}

type quarantiningGatewayReconciler struct {
	*gatewayTestConnectionAccessor
	fail atomic.Bool
}

func (r *quarantiningGatewayReconciler) RepairServer(
	_ context.Context,
	current profile.AgentTeamServer,
) error {
	if r.fail.Load() {
		return errors.New("runtime unavailable")
	}
	r.mu.Lock()
	if connection, ok := r.connections[current.ID].(*gatewayTestConnection); ok {
		connection.serverURL = current.ServerURL
	}
	r.mu.Unlock()
	return nil
}

func (r *quarantiningGatewayReconciler) Reconcile(
	_ context.Context,
	_ *profile.AgentTeamServer,
	_ *profile.AgentTeamServer,
) error {
	if r.fail.Load() {
		return errors.New("runtime unavailable")
	}
	return nil
}

func newGatewayTestRouter(
	t *testing.T,
	upstreams map[string]*httptest.Server,
) (http.Handler, map[string]*gatewayTestConnection) {
	t.Helper()
	dir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, dir)
	connections := make(map[string]*gatewayTestConnection, len(upstreams))
	accessor := &gatewayTestConnectionAccessor{connections: make(map[string]types.ServerConnection)}
	for serverID, upstream := range upstreams {
		server := profile.AgentTeamServer{
			ID:          serverID,
			Enabled:     true,
			Name:        serverID,
			ServerURL:   upstream.URL,
			WorkspaceID: "wsp-" + serverID,
			MachineID:   "mac-" + serverID,
		}
		require.NoError(t, profile.New(dir).Servers().Add(&server))
		connection := &gatewayTestConnection{
			id:         serverID,
			serverURL:  upstream.URL,
			userToken:  "user-token-" + serverID,
			httpClient: upstream.Client(),
		}
		connections[serverID] = connection
		accessor.connections[serverID] = connection
	}
	return newGatewayRuntimeTestRouter(t, accessor), connections
}

func TestAuthoritativeWorkspaceHeaders(t *testing.T) {
	source := make(http.Header)
	source.Set("X-Workspace-Id", "forged")
	source.Set("X-AgentTeam-Route-Workspace-Id", "forged")
	source.Set("X-AgentTeam-Frontend-Build", "forged")
	source.Set("X-AgentTeam-Route-Source", "forged")
	target := cloneGatewayRequestHeaders(source)
	request := httptest.NewRequest(http.MethodGet, "/api/v1/kb", nil)
	request.Header = source
	applyAuthoritativeWorkspaceHeadersForRequest(target, request, "wsp-authoritative")
	require.Equal(t, "wsp-authoritative", target.Get("X-Workspace-Id"))
	require.Equal(t, "wsp-authoritative", target.Get("X-AgentTeam-Route-Workspace-Id"))
	require.Empty(t, target.Get("X-AgentTeam-Frontend-Build"))
	require.Empty(t, target.Get("X-AgentTeam-Route-Source"))

	centralSource := make(http.Header)
	centralTarget := cloneGatewayRequestHeaders(centralSource)
	centralRequest := httptest.NewRequest(http.MethodGet, "/api/v1/auth/session", nil)
	centralRequest.Header = centralSource
	applyAuthoritativeWorkspaceHeadersForRequest(centralTarget, centralRequest, "wsp-authoritative")
	require.Empty(t, centralTarget.Get("X-Workspace-Id"))
	require.Empty(t, centralTarget.Get("X-AgentTeam-Route-Workspace-Id"))

	businessSource := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	businessSource.Header.Set("X-Workspace-Id", "forged")
	businessTarget := cloneGatewayRequestHeaders(businessSource.Header)
	applyAuthoritativeWorkspaceHeadersForRequest(businessTarget, businessSource, "wsp-authoritative")
	require.Empty(t, businessTarget.Get("X-Workspace-Id"))
	require.Empty(t, businessTarget.Get("X-AgentTeam-Route-Workspace-Id"))

	billingSource := httptest.NewRequest(http.MethodGet, "/api/v1/billing/summary", nil)
	billingTarget := cloneGatewayRequestHeaders(billingSource.Header)
	applyAuthoritativeWorkspaceHeadersForRequest(billingTarget, billingSource, "wsp-authoritative")
	require.Equal(t, "wsp-authoritative", billingTarget.Get("X-Workspace-Id"))
	require.Equal(t, "wsp-authoritative", billingTarget.Get("X-AgentTeam-Route-Workspace-Id"))

	pathSource := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp-authoritative/machines", nil)
	pathTarget := cloneGatewayRequestHeaders(pathSource.Header)
	applyAuthoritativeWorkspaceHeadersForRequest(pathTarget, pathSource, "wsp-authoritative")
	require.Equal(t, "wsp-authoritative", pathTarget.Get("X-Workspace-Id"))
	require.Equal(t, "wsp-authoritative", pathTarget.Get("X-AgentTeam-Route-Workspace-Id"))

	pathSource = httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/forged/machines", nil)
	pathSource.Header.Set("X-Workspace-Id", "forged")
	pathSource.Header.Set("X-AgentTeam-Route-Workspace-Id", "forged")
	pathTarget = cloneGatewayRequestHeaders(pathSource.Header)
	applyAuthoritativeWorkspaceHeadersForRequest(pathTarget, pathSource, "wsp-authoritative")
	require.Empty(t, pathTarget.Get("X-Workspace-Id"))
	require.Empty(t, pathTarget.Get("X-AgentTeam-Route-Workspace-Id"))
}

func TestAuthoritativeWorkspaceHeadersKeepGlobalAGSEndpointsUnscoped(t *testing.T) {
	for _, path := range []string{
		"/api/v3",
		"/api/v3/",
		"/api/v3/meta",
		"/api/graphql",
	} {
		t.Run(path, func(t *testing.T) {
			source := httptest.NewRequest(http.MethodGet, path, nil)
			source.Header.Set("X-Workspace-Id", "forged")
			source.Header.Set("X-AgentTeam-Route-Workspace-Id", "forged")
			target := cloneGatewayRequestHeaders(source.Header)

			applyAuthoritativeWorkspaceHeadersForRequest(
				target,
				source,
				"wsp-authoritative",
			)

			require.Empty(t, target.Get("X-Workspace-Id"))
			require.Empty(t, target.Get("X-AgentTeam-Route-Workspace-Id"))
		})
	}
}

func TestAuthoritativeWorkspaceHeadersScopeAGSWorkspaceOperations(t *testing.T) {
	for _, path := range []string{
		"/api/v3/repos/pingcap/agentteam/branches",
		"/api/v3/repos/pingcap/agentteam/tags",
		"/api/v3/repos/pingcap/agentteam/contents/README.md",
		"/api/v3/user",
		"/api/v3/user/repos",
		"/api/v3/search/code",
		"/api/v3/search/labels",
		"/api/ext/v1/user/tokens",
	} {
		t.Run(path, func(t *testing.T) {
			source := httptest.NewRequest(http.MethodGet, path, nil)
			source.Header.Set("X-Workspace-Id", "forged")
			source.Header.Set("X-AgentTeam-Route-Workspace-Id", "forged")
			target := cloneGatewayRequestHeaders(source.Header)
			applyAuthoritativeWorkspaceHeadersForRequest(target, source, "wsp-authoritative")
			require.Equal(t, "wsp-authoritative", target.Get("X-Workspace-Id"))
			require.Equal(t, "wsp-authoritative", target.Get("X-AgentTeam-Route-Workspace-Id"))
		})
	}
}

func TestKnowledgeQAWorkspaceHeaders(t *testing.T) {
	for _, path := range []string{
		"/api/v1/agent-thread/session-channel?subsystem=kb",
		"/api/v1/workspaces/wsp-authoritative/scopes/chn-kb/messages",
		"/api/v1/workspaces/wsp-authoritative/knowledge-qa/policy",
		"/api/v1/workspaces/wsp-authoritative/knowledge-qa/sessions",
		"/api/v1/workspaces/wsp-authoritative/knowledge-qa/sessions/kqs-1/turns",
		"/api/v1/workspaces/wsp-authoritative/knowledge-qa/sessions/kqs-1/turns/kqt-1/retry",
		"/api/v1/workspaces/wsp-authoritative/knowledge-qa/sessions/kqs-1/turns/kqt-1/stop",
		"/api/v1/workspaces/wsp-authoritative/knowledge-qa/sessions/kqs-1/citations/kqc-1",
		"/api/v1/workspaces/wsp-authoritative/knowledge-qa/threads/thr-1/turns",
		"/api/v1/workspaces/wsp-authoritative/knowledge-qa/threads/thr-1",
		"/api/v1/workspaces/wsp-authoritative/threads?subsystem=kb",
		"/api/v1/workspaces/wsp-authoritative/conversations?kinds=replyThread",
	} {
		t.Run(path, func(t *testing.T) {
			source := httptest.NewRequest(http.MethodGet, path, nil)
			source.Header.Set("X-Workspace-Id", "forged")
			source.Header.Set("X-AgentTeam-Route-Workspace-Id", "forged")
			target := cloneGatewayRequestHeaders(source.Header)

			applyAuthoritativeWorkspaceHeadersForRequest(
				target,
				source,
				"wsp-authoritative",
			)

			require.Equal(t, "wsp-authoritative", target.Get("X-Workspace-Id"))
			require.Equal(t, "wsp-authoritative", target.Get("X-AgentTeam-Route-Workspace-Id"))
		})
	}

	unsupported := httptest.NewRequest(
		http.MethodGet,
		"/api/v1/agent-thread/session-channel/extra?subsystem=kb",
		nil,
	)
	unsupported.Header.Set("X-Workspace-Id", "forged")
	target := cloneGatewayRequestHeaders(unsupported.Header)
	applyAuthoritativeWorkspaceHeadersForRequest(
		target,
		unsupported,
		"wsp-authoritative",
	)
	require.Empty(t, target.Get("X-Workspace-Id"))
	require.Empty(t, target.Get("X-AgentTeam-Route-Workspace-Id"))

	missingWorkspace := httptest.NewRequest(
		http.MethodGet,
		"/api/v1/agent-thread/session-channel?subsystem=kb",
		nil,
	)
	missingWorkspace.Header.Set("X-Workspace-Id", "forged")
	missingWorkspace.Header.Set("X-AgentTeam-Route-Workspace-Id", "forged")
	target = cloneGatewayRequestHeaders(missingWorkspace.Header)
	applyAuthoritativeWorkspaceHeadersForRequest(target, missingWorkspace, "")
	require.Empty(t, target.Get("X-Workspace-Id"))
	require.Empty(t, target.Get("X-AgentTeam-Route-Workspace-Id"))
}

func TestWorkspaceEventsStreamWorkspaceID(t *testing.T) {
	require.Equal(t, "wsp-one", workspaceEventsStreamWorkspaceID("/api/v1/workspaces/wsp-one/events/stream"))
	source := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp-one/events/stream", nil)
	header := cloneWebSocketRequestHeader(source.Header)
	applyAuthoritativeWorkspaceHeadersForRequest(header, source, "wsp-one")
	if workspaceEventsStreamWorkspaceID(source.URL.Path) == "wsp-one" {
		header.Set("X-Workspace-Id", "wsp-one")
		header.Set("X-AgentTeam-Route-Workspace-Id", "wsp-one")
	}
	require.Equal(t, "wsp-one", header.Get("X-AgentTeam-Route-Workspace-Id"))

	mismatch := cloneWebSocketRequestHeader(source.Header)
	applyAuthoritativeWorkspaceHeadersForRequest(mismatch, source, "wsp-two")
	if workspaceEventsStreamWorkspaceID(source.URL.Path) == "wsp-two" {
		mismatch.Set("X-AgentTeam-Route-Workspace-Id", "wsp-two")
	}
	require.Empty(t, mismatch.Get("X-AgentTeam-Route-Workspace-Id"))
	require.Empty(t, workspaceEventsStreamWorkspaceID("/api/v1/workspaces/wsp-two/events/stream/extra"))
	require.Empty(t, workspaceEventsStreamWorkspaceID("/api/v1/workspaces/wsp-one/members"))
	require.Empty(t, workspaceEventsStreamWorkspaceID("/api/v1/workspaces/a/b/events/stream"))
}

func newGatewayRuntimeTestRouter(
	t *testing.T,
	accessor types.ServerConnectionAccessor,
) http.Handler {
	t.Helper()
	uiSessions := uisession.NewStore(uisession.Options{})
	servers := profile.Default().Servers()
	require.NoError(t, servers.PrepareRuntimeMode())
	require.NoError(t, servers.ActivateRuntimeMode(
		NewServerRuntimeParticipant(accessor, uiSessions),
	))
	t.Cleanup(servers.DeactivateRuntimeMode)
	return newTestRouterWithUISessions(
		accessor,
		routerTestSessionStore{},
		routerTestTrajectoryStore{},
		uiSessions,
		mcp.NewMessageService(),
	)
}

func newGatewayRuntimeTestRouterWithRegistry(
	t *testing.T,
	accessor types.ServerConnectionAccessor,
	registry gatewayOperationRegistry,
) http.Handler {
	t.Helper()
	uiSessions := uisession.NewStore(uisession.Options{})
	servers := profile.Default().Servers()
	require.NoError(t, servers.PrepareRuntimeMode())
	require.NoError(t, servers.ActivateRuntimeMode(
		NewServerRuntimeParticipant(accessor, uiSessions),
	))
	t.Cleanup(servers.DeactivateRuntimeMode)
	return newTestRouterWithUISessionsAndRegistry(
		accessor,
		routerTestSessionStore{},
		routerTestTrajectoryStore{},
		uiSessions,
		registry,
		mcp.NewMessageService(),
	)
}

// registerBarrierRegistry wraps the real stream registry and lets tests park
// immediately before the delegate Register call (gap B) or immediately after
// the delegate has admitted an Entry but before Register returns (gap C).
// Production code sees only gatewayOperationRegistry; no test hook is added to
// streamregistry.Registry itself.
type registerBarrierRegistry struct {
	delegate *streamregistry.Registry

	beforeDelegate chan struct{}
	releaseBefore  chan struct{}
	afterDelegate  chan struct{}
	releaseAfter   chan struct{}
	registered     chan *streamregistry.Entry
	onRegistered   func(*streamregistry.Entry)

	beforeOnce sync.Once
	afterOnce  sync.Once
}

func (r *registerBarrierRegistry) Register(
	parent context.Context,
	sessionID string,
	generation uint64,
	serverID string,
) (*streamregistry.Entry, error) {
	if r.beforeDelegate != nil {
		r.beforeOnce.Do(func() { close(r.beforeDelegate) })
		<-r.releaseBefore
	}
	entry, err := r.delegate.Register(parent, sessionID, generation, serverID)
	if err == nil && r.registered != nil {
		r.registered <- entry
	}
	if err == nil && r.onRegistered != nil {
		r.onRegistered(entry)
	}
	if r.afterDelegate != nil {
		r.afterOnce.Do(func() { close(r.afterDelegate) })
		<-r.releaseAfter
	}
	return entry, err
}

func newGatewayTestRouterWithRegistry(
	t *testing.T,
	upstreams map[string]*httptest.Server,
) (http.Handler, map[string]*gatewayTestConnection, *streamregistry.Registry) {
	t.Helper()
	dir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, dir)
	connections := make(map[string]*gatewayTestConnection, len(upstreams))
	accessor := &gatewayTestConnectionAccessor{connections: make(map[string]types.ServerConnection)}
	for serverID, upstream := range upstreams {
		server := profile.AgentTeamServer{
			ID:          serverID,
			Enabled:     true,
			Name:        serverID,
			ServerURL:   upstream.URL,
			WorkspaceID: "wsp-" + serverID,
			MachineID:   "mac-" + serverID,
		}
		require.NoError(t, profile.New(dir).Servers().Add(&server))
		connection := &gatewayTestConnection{
			id:         serverID,
			serverURL:  upstream.URL,
			userToken:  "user-token-" + serverID,
			httpClient: upstream.Client(),
		}
		connections[serverID] = connection
		accessor.connections[serverID] = connection
	}
	registry := streamregistry.New()
	return newGatewayRuntimeTestRouterWithRegistry(t, accessor, registry), connections, registry
}

// newGatewayTestRouterWithStore is the A/B/C-barrier test variant: it
// returns the uisession.Store so tests can drive deterministic generation
// mutations (BeginServerTransition, AdvanceServerGeneration) and observe
// the gateway's barrier responses. The store is the SAME instance the
// router installed; mutating it during a request is the deterministic
// equivalent of winning the Get-vs-WatchGeneration race.
func newGatewayTestRouterWithStore(
	t *testing.T,
	upstreams map[string]*httptest.Server,
) (http.Handler, map[string]*gatewayTestConnection, *streamregistry.Registry, *uisession.Store) {
	registry := streamregistry.New()
	router, connections, uiSessions := newGatewayTestRouterWithStoreAndRegistry(t, upstreams, registry)
	return router, connections, registry, uiSessions
}

func newGatewayTestRouterWithStoreAndRegistry(
	t *testing.T,
	upstreams map[string]*httptest.Server,
	registry gatewayOperationRegistry,
) (http.Handler, map[string]*gatewayTestConnection, *uisession.Store) {
	t.Helper()
	dir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, dir)
	connections := make(map[string]*gatewayTestConnection, len(upstreams))
	accessor := &gatewayTestConnectionAccessor{connections: make(map[string]types.ServerConnection)}
	for serverID, upstream := range upstreams {
		server := profile.AgentTeamServer{
			ID:          serverID,
			Enabled:     true,
			Name:        serverID,
			ServerURL:   upstream.URL,
			WorkspaceID: "wsp-" + serverID,
			MachineID:   "mac-" + serverID,
		}
		require.NoError(t, profile.New(dir).Servers().Add(&server))
		connection := &gatewayTestConnection{
			id:         serverID,
			serverURL:  upstream.URL,
			userToken:  "user-token-" + serverID,
			httpClient: upstream.Client(),
		}
		connections[serverID] = connection
		accessor.connections[serverID] = connection
	}
	uiSessions := uisession.NewStore(uisession.Options{})
	servers := profile.Default().Servers()
	require.NoError(t, servers.PrepareRuntimeMode())
	require.NoError(t, servers.ActivateRuntimeMode(
		NewServerRuntimeParticipant(accessor, uiSessions),
	))
	t.Cleanup(servers.DeactivateRuntimeMode)
	router := newTestRouterWithUISessionsAndRegistry(
		accessor,
		routerTestSessionStore{},
		routerTestTrajectoryStore{},
		uiSessions,
		registry,
		mcp.NewMessageService(),
	)
	return router, connections, uiSessions
}

func newBareGatewayTestHandlerWithRegistry(
	t *testing.T,
	upstreams map[string]*httptest.Server,
	registry gatewayOperationRegistry,
) (http.Handler, map[string]*gatewayTestConnection, *uisession.Store) {
	t.Helper()
	_, connections, uiSessions := newGatewayTestRouterWithStoreAndRegistry(t, upstreams, registry)
	accessor := &gatewayTestConnectionAccessor{connections: make(map[string]types.ServerConnection, len(connections))}
	for serverID, connection := range connections {
		accessor.connections[serverID] = connection
	}
	return newBareGatewayHandler(accessor, uiSessions, registry), connections, uiSessions
}

func selectGatewayServer(
	t *testing.T,
	router http.Handler,
	sessionToken string,
	serverID string,
) {
	t.Helper()
	req := httptest.NewRequest(
		http.MethodPut,
		"/api/internal/servers/active",
		strings.NewReader(`{"serverId":"`+serverID+`"}`),
	)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)
	require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
}

func selectGatewayStore(
	t *testing.T,
	store *uisession.Store,
	sessionToken string,
	serverID string,
) {
	t.Helper()
	expiresAt, ok := token.SessionExpiresAt(sessionToken)
	require.True(t, ok)
	_, err := store.Select(sessionToken, expiresAt, serverID)
	require.NoError(t, err)
}

func TestGatewayProxiesArbitraryPathQueryMethodBodyAndSanitizesHeaders(t *testing.T) {
	type seenRequest struct {
		method           string
		path             string
		query            string
		body             string
		auth             string
		cookie           string
		forwarded        string
		host             string
		hop              string
		origin           string
		referer          string
		xForwardedFor    string
		xForwardedHost   string
		xRealIP          string
		workspaceID      string
		routeWorkspaceID string
	}
	seen := make(chan seenRequest, 1)
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		body, _ := io.ReadAll(request.Body)
		seen <- seenRequest{
			method:           request.Method,
			path:             request.URL.EscapedPath(),
			query:            request.URL.RawQuery,
			body:             string(body),
			auth:             request.Header.Get("Authorization"),
			cookie:           request.Header.Get("Cookie"),
			forwarded:        request.Header.Get("Forwarded"),
			host:             request.Host,
			hop:              request.Header.Get("X-Hop"),
			origin:           request.Header.Get("Origin"),
			referer:          request.Header.Get("Referer"),
			xForwardedFor:    request.Header.Get("X-Forwarded-For"),
			xForwardedHost:   request.Header.Get("X-Forwarded-Host"),
			xRealIP:          request.Header.Get("X-Real-Ip"),
			workspaceID:      request.Header.Get("X-Workspace-Id"),
			routeWorkspaceID: request.Header.Get("X-AgentTeam-Route-Workspace-Id"),
		}
		writer.Header().Set("X-Upstream", "ok")
		writer.WriteHeader(http.StatusCreated)
		_, _ = writer.Write([]byte("proxied"))
	}))
	defer upstream.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(
		http.MethodPost,
		"/api/v3/repos/a%2Fb?limit=2&cursor=x",
		strings.NewReader("upload-body"),
	)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	req.Header.Set("Cookie", "local=secret")
	req.Header.Set("Connection", "X-Hop")
	req.Header.Set("X-Hop", "strip-me")
	req.Header.Set("Forwarded", "for=192.0.2.1;host=attacker.invalid")
	req.Header.Set("X-Forwarded-For", "192.0.2.1")
	req.Header.Set("X-Forwarded-Host", "attacker.invalid")
	req.Header.Set("X-Real-Ip", "192.0.2.1")
	req.Header.Set("Origin", "https://attacker.invalid")
	req.Header.Set("Referer", "https://attacker.invalid/phish")
	req.Header.Set("X-Workspace-Id", "forged")
	req.Header.Set("X-AgentTeam-Route-Workspace-Id", "forged")
	req.Host = "attacker.invalid"
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusCreated, rec.Code, rec.Body.String())
	require.Equal(t, "proxied", rec.Body.String())
	require.Equal(t, "ok", rec.Header().Get("X-Upstream"))
	got := <-seen
	require.Equal(t, http.MethodPost, got.method)
	require.Equal(t, "/api/v3/repos/a%2Fb", got.path)
	require.Equal(t, "limit=2&cursor=x", got.query)
	require.Equal(t, "upload-body", got.body)
	require.Equal(t, "Bearer user-token-server-a", got.auth)
	require.Empty(t, got.cookie)
	require.Empty(t, got.forwarded)
	require.Empty(t, got.hop)
	require.Empty(t, got.origin)
	require.Empty(t, got.referer)
	require.Empty(t, got.xForwardedFor)
	require.Empty(t, got.xForwardedHost)
	require.Empty(t, got.xRealIP)
	require.Empty(t, got.workspaceID)
	require.Empty(t, got.routeWorkspaceID)
	upstreamURL, err := url.Parse(upstream.URL)
	require.NoError(t, err)
	require.Equal(t, upstreamURL.Host, got.host)
}

func TestGatewayReturnsJSON404ForLocalAndUnknownPaths(t *testing.T) {
	var upstreamCalls atomic.Int32
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		upstreamCalls.Add(1)
		_, _ = writer.Write([]byte("remote:" + request.URL.Path))
	}))
	defer upstream.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	for _, path := range []string{
		"/api/internal/does-not-exist",
		"/api/does-not-exist",
		"/agentteam-logo-black.svg",
		"/assets/not-in-manifest.js",
		"/workspaces/wsp-1/machines",
	} {
		t.Run(path, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, path, nil)
			req.Header.Set("Authorization", "Bearer "+sessionToken)
			req.Header.Set("Accept", "text/html,application/xhtml+xml")
			req.Header.Set("Sec-Fetch-Mode", "navigate")
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)

			require.Equal(t, http.StatusNotFound, rec.Code, rec.Body.String())
			require.Contains(t, rec.Header().Get("Content-Type"), "application/json")
			require.JSONEq(t, `{"error":"Not Found"}`, rec.Body.String())
		})
	}
	require.Equal(t, int32(0), upstreamCalls.Load())
}

func TestGatewayKeepsSelectedServerCatchAllForCurrentAndFutureRoutes(t *testing.T) {
	type seenRequest struct {
		path  string
		query string
	}
	seen := make(chan seenRequest, 5)
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		seen <- seenRequest{path: request.URL.EscapedPath(), query: request.URL.RawQuery}
		_, _ = writer.Write([]byte("remote:" + request.URL.EscapedPath()))
	}))
	defer upstream.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	for _, path := range []string{
		"/api/v1/me",
		"/api/v3/repos/pingcap/agentteam",
		"/api/ext/v1/user/orgs",
		"/api/graphql",
		"/api/attachments/att%2F1",
		"/meshx/future/repos/pingcap/agentteam",
	} {
		t.Run(path, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, path, nil)
			req.Header.Set("Authorization", "Bearer "+sessionToken)
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)

			require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
			require.Equal(t, "remote:"+path, rec.Body.String())
			got := <-seen
			require.Equal(t, path, got.path)
			require.Empty(t, got.query)
		})
	}
}

func TestGatewayKeepsSelectionsIsolatedAcrossDesktopSessions(t *testing.T) {
	serverA := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		_, _ = writer.Write([]byte("server-a"))
	}))
	defer serverA.Close()
	serverB := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		_, _ = writer.Write([]byte("server-b"))
	}))
	defer serverB.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{
		"server-a": serverA,
		"server-b": serverB,
	})
	windowA := token.CreateSession()
	windowB := token.CreateSession()
	selectGatewayServer(t, router, windowA, "server-a")
	selectGatewayServer(t, router, windowB, "server-b")

	for sessionToken, want := range map[string]string{
		windowA: "server-a",
		windowB: "server-b",
	} {
		req := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
		req.Header.Set("Authorization", "Bearer "+sessionToken)
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)
		require.Equal(t, http.StatusOK, rec.Code)
		require.Equal(t, want, rec.Body.String())
	}
}

func TestGatewayReturnsStableErrorsForMissingSelectionAndUnreachableRemote(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	upstream.Close()
	sessionToken := token.CreateSession()

	missingReq := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	missingReq.Header.Set("Authorization", "Bearer "+sessionToken)
	missingRec := httptest.NewRecorder()
	router.ServeHTTP(missingRec, missingReq)
	require.Equal(t, http.StatusConflict, missingRec.Code)
	require.Contains(t, missingRec.Body.String(), `"code":"ACTIVE_SERVER_REQUIRED"`)

	selectGatewayServer(t, router, sessionToken, "server-a")
	unreachableReq := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	unreachableReq.Header.Set("Authorization", "Bearer "+sessionToken)
	unreachableRec := httptest.NewRecorder()
	router.ServeHTTP(unreachableRec, unreachableReq)
	require.Equal(t, http.StatusBadGateway, unreachableRec.Code)
	require.Contains(t, unreachableRec.Body.String(), `"code":"REMOTE_UNREACHABLE"`)
}

func TestGatewayRefreshesUserTokenOnceAfterUnauthorized(t *testing.T) {
	var requests atomic.Int32
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		requests.Add(1)
		if request.Header.Get("Authorization") == "Bearer fresh-token" {
			_, _ = writer.Write([]byte("ok"))
			return
		}
		writer.WriteHeader(http.StatusUnauthorized)
	}))
	defer upstream.Close()
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	connection := connections["server-a"]
	connection.userToken = "stale-token"
	connection.refreshToken = "fresh-token"
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
	require.Equal(t, "ok", rec.Body.String())
	require.Equal(t, int32(2), requests.Load())
	require.Equal(t, int32(1), connection.invalidates.Load())
}

func TestGatewayCoalescesConcurrentUnauthorizedTokenRefresh(t *testing.T) {
	var requests atomic.Int32
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		requests.Add(1)
		if request.Header.Get("Authorization") == "Bearer fresh-token" {
			_, _ = writer.Write([]byte("ok"))
			return
		}
		writer.WriteHeader(http.StatusUnauthorized)
	}))
	defer upstream.Close()
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	connection := connections["server-a"]
	connection.userToken = "stale-token"
	connection.refreshToken = "fresh-token"
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	const callers = 16
	var wg sync.WaitGroup
	failures := make(chan string, callers)
	for range callers {
		wg.Add(1)
		go func() {
			defer wg.Done()
			req := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
			req.Header.Set("Authorization", "Bearer "+sessionToken)
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)
			if rec.Code != http.StatusOK || rec.Body.String() != "ok" {
				failures <- rec.Body.String()
			}
		}()
	}
	wg.Wait()
	close(failures)
	for failure := range failures {
		t.Errorf("gateway request failed: %s", failure)
	}
	require.Equal(t, int32(1), connection.invalidates.Load())
	require.GreaterOrEqual(t, requests.Load(), int32(callers))
}

func TestGatewayDropsLateResponseAfterActiveServerGenerationChanges(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	upstreamA := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		close(started)
		<-release
		_, _ = writer.Write([]byte("stale"))
	}))
	defer upstreamA.Close()
	upstreamB := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		_, _ = writer.Write([]byte("fresh"))
	}))
	defer upstreamB.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{
		"server-a": upstreamA,
		"server-b": upstreamB,
	})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	result := make(chan *httptest.ResponseRecorder, 1)
	go func() {
		req := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
		req.Header.Set("Authorization", "Bearer "+sessionToken)
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)
		result <- rec
	}()
	select {
	case <-started:
	case <-time.After(3 * time.Second):
		t.Fatal("timed out waiting for first upstream")
	}
	selectGatewayServer(t, router, sessionToken, "server-b")
	close(release)

	rec := <-result
	require.Equal(t, http.StatusConflict, rec.Code, rec.Body.String())
	require.NotContains(t, rec.Body.String(), "stale")
	require.Contains(t, rec.Body.String(), `"code":"ACTIVE_SERVER_CHANGED"`)
}

func TestGatewayRejectsRequestsDuringServerIdentityTransition(t *testing.T) {
	oldUpstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		_, _ = writer.Write([]byte("old"))
	}))
	defer oldUpstream.Close()
	newUpstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		_, _ = writer.Write([]byte("new"))
	}))
	defer newUpstream.Close()

	dir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, dir)
	server := profile.AgentTeamServer{
		ID:          "server-a",
		Enabled:     true,
		Name:        "Production",
		ServerURL:   oldUpstream.URL,
		APIKey:      "old-key",
		WorkspaceID: "wsp-old",
		MachineID:   "mac-old",
	}
	require.NoError(t, profile.New(dir).Servers().Add(&server))
	oldConnection := &gatewayTestConnection{
		id:         server.ID,
		serverURL:  oldUpstream.URL,
		userToken:  "old-user-token",
		httpClient: oldUpstream.Client(),
	}
	newConnection := &gatewayTestConnection{
		id:         server.ID,
		serverURL:  newUpstream.URL,
		userToken:  "new-user-token",
		httpClient: newUpstream.Client(),
	}
	accessor := &blockingGatewayReconciler{
		gatewayTestConnectionAccessor: &gatewayTestConnectionAccessor{
			connections: map[string]types.ServerConnection{server.ID: oldConnection},
		},
		entered:       make(chan struct{}),
		release:       make(chan struct{}),
		newConnection: newConnection,
	}
	router := newGatewayRuntimeTestRouter(t, accessor)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, server.ID)

	updateResult := make(chan *httptest.ResponseRecorder, 1)
	go func() {
		request := httptest.NewRequest(
			http.MethodPut,
			"/api/internal/servers/server-a",
			strings.NewReader(`{
				"enabled": true,
				"config": {
					"server_url": "`+newUpstream.URL+`",
					"api_key": "new-key",
					"workspace_id": "wsp-new",
					"machine_id": "mac-new"
				}
			}`),
		)
		request.Header.Set("Authorization", "Bearer "+sessionToken)
		request.Header.Set("Content-Type", "application/json")
		recorder := httptest.NewRecorder()
		router.ServeHTTP(recorder, request)
		updateResult <- recorder
	}()
	select {
	case <-accessor.entered:
	case <-time.After(3 * time.Second):
		t.Fatal("timed out waiting for blocked server reconcile")
	}

	duringReq := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	duringReq.Header.Set("Authorization", "Bearer "+sessionToken)
	duringRec := httptest.NewRecorder()
	router.ServeHTTP(duringRec, duringReq)
	require.Equal(t, http.StatusConflict, duringRec.Code, duringRec.Body.String())
	require.Contains(t, duringRec.Body.String(), `"code":"ACTIVE_SERVER_CHANGED"`)
	require.NotContains(t, duringRec.Body.String(), "old")
	require.NotContains(t, duringRec.Body.String(), "new")

	close(accessor.release)
	updateRec := <-updateResult
	require.Equal(t, http.StatusOK, updateRec.Code, updateRec.Body.String())

	afterReq := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	afterReq.Header.Set("Authorization", "Bearer "+sessionToken)
	afterRec := httptest.NewRecorder()
	router.ServeHTTP(afterRec, afterReq)
	require.Equal(t, http.StatusOK, afterRec.Code, afterRec.Body.String())
	require.Equal(t, "new", afterRec.Body.String())
}

func TestGatewayReopensOldIdentityAfterServerReconcileRollback(t *testing.T) {
	oldUpstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		_, _ = writer.Write([]byte("old"))
	}))
	defer oldUpstream.Close()

	dir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, dir)
	server := profile.AgentTeamServer{
		ID:          "server-a",
		Enabled:     true,
		Name:        "Production",
		ServerURL:   oldUpstream.URL,
		APIKey:      "old-key",
		WorkspaceID: "wsp-old",
		MachineID:   "mac-old",
	}
	require.NoError(t, profile.New(dir).Servers().Add(&server))
	oldConnection := &gatewayTestConnection{
		id:         server.ID,
		serverURL:  oldUpstream.URL,
		userToken:  "old-user-token",
		httpClient: oldUpstream.Client(),
	}
	accessor := &failingGatewayReconciler{
		gatewayTestConnectionAccessor: &gatewayTestConnectionAccessor{
			connections: map[string]types.ServerConnection{server.ID: oldConnection},
		},
	}
	router := newGatewayRuntimeTestRouter(t, accessor)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, server.ID)

	updateReq := httptest.NewRequest(
		http.MethodPut,
		"/api/internal/servers/server-a",
		strings.NewReader(`{
			"enabled": true,
			"config": {
				"server_url": "https://new.example.test",
				"api_key": "new-key",
				"workspace_id": "wsp-new",
				"machine_id": "mac-new"
			}
		}`),
	)
	updateReq.Header.Set("Authorization", "Bearer "+sessionToken)
	updateReq.Header.Set("Content-Type", "application/json")
	updateRec := httptest.NewRecorder()
	router.ServeHTTP(updateRec, updateReq)
	require.Equal(t, http.StatusInternalServerError, updateRec.Code, updateRec.Body.String())

	activeReq := httptest.NewRequest(http.MethodGet, "/api/internal/servers/active", nil)
	activeReq.Header.Set("Authorization", "Bearer "+sessionToken)
	activeRec := httptest.NewRecorder()
	router.ServeHTTP(activeRec, activeReq)
	require.Equal(t, http.StatusOK, activeRec.Code, activeRec.Body.String())
	require.Contains(t, activeRec.Body.String(), `"generation":3`)
	require.Contains(t, activeRec.Body.String(), `"workspaceId":"wsp-old"`)
	require.Contains(t, activeRec.Body.String(), `"machineId":"mac-old"`)

	afterReq := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	afterReq.Header.Set("Authorization", "Bearer "+sessionToken)
	afterRec := httptest.NewRecorder()
	router.ServeHTTP(afterRec, afterReq)
	require.Equal(t, http.StatusOK, afterRec.Code, afterRec.Body.String())
	require.Equal(t, "old", afterRec.Body.String())
}

func TestGatewayRuntimeRollbackFailureQuarantinesAllSessionsUntilRepair(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		_, _ = writer.Write([]byte("old"))
	}))
	defer upstream.Close()

	dir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, dir)
	server := profile.AgentTeamServer{
		ID:          "server-a",
		Enabled:     true,
		Name:        "Production",
		ServerURL:   upstream.URL,
		APIKey:      "old-key",
		WorkspaceID: "wsp-old",
		MachineID:   "mac-old",
	}
	require.NoError(t, profile.New(dir).Servers().Add(&server))
	connection := &gatewayTestConnection{
		id:         server.ID,
		serverURL:  upstream.URL,
		userToken:  "old-user-token",
		httpClient: upstream.Client(),
	}
	accessor := &quarantiningGatewayReconciler{
		gatewayTestConnectionAccessor: &gatewayTestConnectionAccessor{
			connections: map[string]types.ServerConnection{server.ID: connection},
		},
	}
	router := newGatewayRuntimeTestRouter(t, accessor)
	windowA := token.CreateSession()
	windowB := token.CreateSession()
	selectGatewayServer(t, router, windowA, server.ID)
	selectGatewayServer(t, router, windowB, server.ID)

	accessor.fail.Store(true)
	updateReq := httptest.NewRequest(
		http.MethodPut,
		"/api/internal/servers/server-a",
		strings.NewReader(`{
			"enabled": true,
			"config": {
				"server_url": "https://new.example.test",
				"api_key": "new-key",
				"workspace_id": "wsp-new",
				"machine_id": "mac-new"
			}
		}`),
	)
	updateReq.Header.Set("Authorization", "Bearer "+windowA)
	updateReq.Header.Set("Content-Type", "application/json")
	updateRec := httptest.NewRecorder()
	router.ServeHTTP(updateRec, updateReq)
	require.Equal(t, http.StatusInternalServerError, updateRec.Code, updateRec.Body.String())
	require.Contains(t, updateRec.Body.String(), "runtime rollback failed")
	require.True(t, profile.Default().Servers().ServerRuntimeHalted(server.ID))

	for _, sessionToken := range []string{windowA, windowB} {
		activeReq := httptest.NewRequest(http.MethodGet, "/api/internal/servers/active", nil)
		activeReq.Header.Set("Authorization", "Bearer "+sessionToken)
		activeRec := httptest.NewRecorder()
		router.ServeHTTP(activeRec, activeReq)
		require.Equal(t, http.StatusConflict, activeRec.Code, activeRec.Body.String())
		require.Contains(t, activeRec.Body.String(), `"code":"ACTIVE_SERVER_REQUIRED"`)

		selectReq := httptest.NewRequest(
			http.MethodPut,
			"/api/internal/servers/active",
			strings.NewReader(`{"serverId":"server-a"}`),
		)
		selectReq.Header.Set("Authorization", "Bearer "+sessionToken)
		selectReq.Header.Set("Content-Type", "application/json")
		selectRec := httptest.NewRecorder()
		router.ServeHTTP(selectRec, selectReq)
		require.Equal(t, http.StatusConflict, selectRec.Code, selectRec.Body.String())
		require.Contains(t, selectRec.Body.String(), `"code":"SERVER_UNAVAILABLE"`)
	}

	blockedReq := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	blockedReq.Header.Set("Authorization", "Bearer "+windowA)
	blockedRec := httptest.NewRecorder()
	router.ServeHTTP(blockedRec, blockedReq)
	require.Equal(t, http.StatusConflict, blockedRec.Code, blockedRec.Body.String())
	require.Contains(t, blockedRec.Body.String(), `"code":"ACTIVE_SERVER_REQUIRED"`)

	accessor.fail.Store(false)
	require.NoError(t, profile.Default().Servers().RepairServerRuntime(server.ID))
	require.False(t, profile.Default().Servers().ServerRuntimeHalted(server.ID))
	selectGatewayServer(t, router, windowA, server.ID)

	repairedReq := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	repairedReq.Header.Set("Authorization", "Bearer "+windowA)
	repairedRec := httptest.NewRecorder()
	router.ServeHTTP(repairedRec, repairedReq)
	require.Equal(t, http.StatusOK, repairedRec.Code, repairedRec.Body.String())
	require.Equal(t, "old", repairedRec.Body.String())
}

func TestGatewayCancelsStreamingHTTPWhenActiveServerGenerationChanges(t *testing.T) {
	headersSent := make(chan struct{})
	requestCanceled := make(chan struct{})
	registry := streamregistry.New()
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.WriteHeader(http.StatusOK)
		writer.(http.Flusher).Flush()
		close(headersSent)
		<-request.Context().Done()
		close(requestCanceled)
	}))
	defer upstream.Close()
	router, _, _ := newGatewayTestRouterWithStoreAndRegistry(
		t,
		map[string]*httptest.Server{"server-a": upstream},
		registry,
	)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	result := make(chan *httptest.ResponseRecorder, 1)
	go func() {
		req := httptest.NewRequest(http.MethodGet, "/api/v1/download", nil)
		req.Header.Set("Authorization", "Bearer "+sessionToken)
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)
		result <- rec
	}()
	select {
	case <-headersSent:
	case <-time.After(3 * time.Second):
		t.Fatal("timed out waiting for upstream response headers")
	}
	selectGatewayServer(t, router, sessionToken, "server-a")

	select {
	case <-requestCanceled:
	case <-time.After(3 * time.Second):
		t.Fatal("generation change did not cancel the in-flight upstream request")
	}
	select {
	case rec := <-result:
		require.NotContains(t, rec.Body.String(), "stale")
	case <-time.After(3 * time.Second):
		t.Fatal("gateway request did not stop after generation changed")
	}
	require.Eventually(t, func() bool {
		return registry.ActiveGatewayOps() == 0
	}, time.Second, 5*time.Millisecond, "generation-canceled stream entry did not release")
}

func TestGatewaySSEGenerationChangeEndsCommittedStreamAndReleasesEntry(t *testing.T) {
	const eventPayloadSize = 8 << 10
	event := "data: " + strings.Repeat("x", eventPayloadSize) + "\n\n"
	firstEventSent := make(chan struct{})
	upstreamCanceled := make(chan struct{})
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "text/event-stream")
		_, _ = writer.Write([]byte(event))
		writer.(http.Flusher).Flush()
		close(firstEventSent)
		<-request.Context().Done()
		close(upstreamCanceled)
	}))
	defer upstream.Close()

	realRegistry := streamregistry.New()
	registered := make(chan *streamregistry.Entry, 1)
	registry := &registerBarrierRegistry{
		delegate:   realRegistry,
		registered: registered,
	}
	router, _, _ := newGatewayTestRouterWithStoreAndRegistry(
		t,
		map[string]*httptest.Server{"server-a": upstream},
		registry,
	)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")
	listener := httptest.NewServer(router)
	defer listener.Close()

	req, err := http.NewRequest(http.MethodGet, listener.URL+"/api/v1/events", nil)
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	resp, err := listener.Client().Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	require.Equal(t, http.StatusOK, resp.StatusCode)

	reader := bufio.NewReader(resp.Body)
	eventLine, err := reader.ReadString('\n')
	require.NoError(t, err)
	separator, err := reader.ReadString('\n')
	require.NoError(t, err)
	require.Equal(t, "data: "+strings.Repeat("x", eventPayloadSize)+"\n\n", eventLine+separator)
	select {
	case <-firstEventSent:
	case <-time.After(time.Second):
		t.Fatal("upstream did not send the first SSE event")
	}
	var entry *streamregistry.Entry
	select {
	case entry = <-registered:
	case <-time.After(time.Second):
		t.Fatal("SSE entry was not captured")
	}
	require.Equal(t, int64(1), realRegistry.ActiveGatewayOps())

	selectGatewayServer(t, router, sessionToken, "server-a")

	select {
	case <-upstreamCanceled:
	case <-time.After(time.Second):
		t.Fatal("generation change did not cancel committed SSE upstream")
	}
	require.Eventually(t, func() bool {
		return errors.Is(context.Cause(entry.Context()), streamregistry.ReasonGenerationChange)
	}, time.Second, 5*time.Millisecond, "SSE entry did not retain typed generation-change cause")
	streamEnded := make(chan error, 1)
	go func() {
		_, readErr := reader.ReadByte()
		streamEnded <- readErr
	}()
	select {
	case readErr := <-streamEnded:
		require.ErrorIs(t, readErr, io.EOF,
			"committed SSE must end instead of synthesizing a new status/body")
	case <-time.After(time.Second):
		t.Fatal("committed SSE did not reach EOF after generation change")
	}
	require.Eventually(t, func() bool {
		return realRegistry.ActiveGatewayOps() == 0
	}, time.Second, 5*time.Millisecond, "SSE entry did not release after generation cancel")
}

func TestGatewayRewritesSameOriginRedirectAndRejectsCrossOriginRedirect(t *testing.T) {
	var upstreamURL string
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		switch request.URL.Path {
		case "/api/v1/redirect/same":
			http.Redirect(writer, request, upstreamURL+"/api/v1/redirect/next?x=1", http.StatusFound)
			return
		case "/api/v1/redirect/userinfo":
			target, err := url.Parse(upstreamURL)
			require.NoError(t, err)
			target.User = url.User("user")
			target.Path = "/api/v1/redirect/next"
			http.Redirect(writer, request, target.String(), http.StatusFound)
			return
		}
		http.Redirect(writer, request, "https://evil.example.test/phish", http.StatusFound)
	}))
	defer upstream.Close()
	upstreamURL = upstream.URL
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	sameReq := httptest.NewRequest(http.MethodGet, "/api/v1/redirect/same", nil)
	sameReq.Header.Set("Authorization", "Bearer "+sessionToken)
	sameReq.Host = "127.0.0.1:8084"
	sameRec := httptest.NewRecorder()
	router.ServeHTTP(sameRec, sameReq)
	require.Equal(t, http.StatusFound, sameRec.Code)
	require.Equal(t, "http://127.0.0.1:8084/api/v1/redirect/next?x=1", sameRec.Header().Get("Location"))

	crossReq := httptest.NewRequest(http.MethodGet, "/api/v1/redirect/cross", nil)
	crossReq.Header.Set("Authorization", "Bearer "+sessionToken)
	crossRec := httptest.NewRecorder()
	router.ServeHTTP(crossRec, crossReq)
	require.Equal(t, http.StatusBadGateway, crossRec.Code)
	require.Contains(t, crossRec.Body.String(), `"code":"REMOTE_REDIRECT_REJECTED"`)

	userinfoReq := httptest.NewRequest(http.MethodGet, "/api/v1/redirect/userinfo", nil)
	userinfoReq.Header.Set("Authorization", "Bearer "+sessionToken)
	userinfoRec := httptest.NewRecorder()
	router.ServeHTTP(userinfoRec, userinfoReq)
	require.Equal(t, http.StatusBadGateway, userinfoRec.Code)
	require.Contains(t, userinfoRec.Body.String(), `"code":"REMOTE_REDIRECT_REJECTED"`)
}

func TestGatewayRejectsRegistryAndConnectionOriginMismatchBeforeDial(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("mismatched upstream must not be dialed")
	}))
	defer upstream.Close()
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	connections["server-a"].serverURL = "https://attacker.invalid"
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusBadGateway, rec.Code)
	require.Contains(t, rec.Body.String(), `"code":"REMOTE_UNREACHABLE"`)
}

func TestGatewayDoesNotRetryLargeUnauthorizedRequestBody(t *testing.T) {
	var requests atomic.Int32
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		requests.Add(1)
		writer.WriteHeader(http.StatusUnauthorized)
	}))
	defer upstream.Close()
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	connection := connections["server-a"]
	connection.userToken = "stale-token"
	connection.refreshToken = "fresh-token"
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(
		http.MethodPost,
		"/api/v1/upload",
		strings.NewReader(strings.Repeat("x", maxRetryableGatewayBodyBytes+1)),
	)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusUnauthorized, rec.Code)
	require.Equal(t, int32(1), requests.Load())
	require.Equal(t, int32(0), connection.invalidates.Load())
}

func TestGatewayCopiesSafeResponseHeadersAndDropsRestrictedHeaders(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		writer.Header().Set("Connection", "X-Upstream-Hop")
		writer.Header().Set("X-Upstream-Hop", "strip-me")
		writer.Header().Set("Set-Cookie", "upstream=secret")
		writer.Header().Set("ETag", `"workflow-plan:pln_1:1"`)
		writer.Header().Set("X-Plan-Draft-ETag", `"workflow-plan-draft:pln_1:1"`)
		writer.Header().Set("Content-Disposition", `attachment; filename="workflow.yaml"`)
		writer.Header().Set("X-Safe", "keep")
		_, _ = writer.Write([]byte("ok"))
	}))
	defer upstream.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Empty(t, rec.Header().Get("Set-Cookie"))
	require.Empty(t, rec.Header().Get("X-Upstream-Hop"))
	require.Empty(t, rec.Header().Get("Connection"))
	require.Equal(t, "keep", rec.Header().Get("X-Safe"))
	require.Equal(t, `"workflow-plan:pln_1:1"`, rec.Header().Get("ETag"))
	require.Equal(t, `"workflow-plan-draft:pln_1:1"`, rec.Header().Get("X-Plan-Draft-ETag"))
	require.Equal(t, `attachment; filename="workflow.yaml"`, rec.Header().Get("Content-Disposition"))
}

func TestGatewayProxiesAttachmentDownloadWithQuerySessionCredential(t *testing.T) {
	seen := make(chan struct {
		path  string
		query url.Values
		auth  string
	}, 1)
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		seen <- struct {
			path  string
			query url.Values
			auth  string
		}{
			path:  request.URL.EscapedPath(),
			query: request.URL.Query(),
			auth:  request.Header.Get("Authorization"),
		}
		writer.Header().Set("Content-Disposition", `attachment; filename="report.txt"`)
		_, _ = writer.Write([]byte("attachment body"))
	}))
	defer upstream.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(
		http.MethodGet,
		"/api/v1/workspaces/wsp%2Fone/attachments/att%201?accessToken="+url.QueryEscape(sessionToken),
		nil,
	)
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8")
	req.Header.Set("Sec-Fetch-Mode", "navigate")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
	require.Equal(t, "attachment body", rec.Body.String())
	require.Equal(t, `attachment; filename="report.txt"`, rec.Header().Get("Content-Disposition"))
	request := <-seen
	require.Equal(t, "/api/v1/workspaces/wsp%2Fone/attachments/att%201", request.path)
	require.Equal(t, []string{"user-token-server-a"}, request.query["accessToken"])
	require.Equal(t, "Bearer user-token-server-a", request.auth)
}

func TestGatewayRewritesDesktopSessionQueryCredentialForRealtimeCompatibility(t *testing.T) {
	seenQuery := make(chan url.Values, 1)
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		seenQuery <- request.URL.Query()
		_, _ = writer.Write([]byte("ok"))
	}))
	defer upstream.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(
		http.MethodGet,
		"/api/v1/events/stream?clientId=client-1&accessToken="+url.QueryEscape(sessionToken),
		nil,
	)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	query := <-seenQuery
	require.Equal(t, "client-1", query.Get("clientId"))
	require.Equal(t, "user-token-server-a", query.Get("accessToken"))
}

func TestGatewayNormalizesQueryCredentialsToOneUpstreamToken(t *testing.T) {
	tests := []struct {
		name      string
		rawQuery  func(string) string
		wantKey   string
		absentKey string
	}{
		{
			name: "duplicate accessToken values and competing token key",
			rawQuery: func(sessionToken string) string {
				return "clientId=client-1&accessToken=client-controlled&accessToken=" +
					url.QueryEscape(sessionToken) +
					"&token=also-client-controlled"
			},
			wantKey:   "accessToken",
			absentKey: "token",
		},
		{
			name: "token authentication removes competing accessToken key",
			rawQuery: func(sessionToken string) string {
				return "clientId=client-1&accessToken=client-controlled&token=" +
					url.QueryEscape(sessionToken) +
					"&token=also-client-controlled"
			},
			wantKey:   "token",
			absentKey: "accessToken",
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			seenQuery := make(chan url.Values, 1)
			upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
				seenQuery <- request.URL.Query()
				_, _ = writer.Write([]byte("ok"))
			}))
			defer upstream.Close()
			router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
			sessionToken := token.CreateSession()
			selectGatewayServer(t, router, sessionToken, "server-a")

			req := httptest.NewRequest(
				http.MethodGet,
				"/api/v1/events/stream?"+test.rawQuery(sessionToken),
				nil,
			)
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)

			require.Equal(t, http.StatusBadRequest, rec.Code)
			require.Contains(t, rec.Body.String(), "INVALID_QUERY_CREDENTIALS")
		})
	}
}

func TestGatewayNormalizesHeaderAuthenticatedQueryCredentials(t *testing.T) {
	seenQuery := make(chan url.Values, 1)
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		seenQuery <- request.URL.Query()
		_, _ = writer.Write([]byte("ok"))
	}))
	defer upstream.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(
		http.MethodGet,
		"/api/v1/events/stream?clientId=client-1"+
			"&accessToken=client-controlled"+
			"&accessToken="+url.QueryEscape(sessionToken)+
			"&token=also-client-controlled",
		nil,
	)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String())
	require.Contains(t, rec.Body.String(), "INVALID_QUERY_CREDENTIALS")
}

func TestGatewayRejectsMalformedCredentialQueryBeforeAnyAuthenticationPath(t *testing.T) {
	tests := []struct {
		name      string
		authorize bool
	}{
		{name: "query session authentication"},
		{name: "header session authentication", authorize: true},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var upstreamCalls atomic.Int32
			upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
				upstreamCalls.Add(1)
			}))
			defer upstream.Close()
			router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
			sessionToken := token.CreateSession()
			selectGatewayServer(t, router, sessionToken, "server-a")

			req := httptest.NewRequest(
				http.MethodGet,
				"/api/v1/events/stream?accessToken="+url.QueryEscape(sessionToken)+"&broken=%zz",
				nil,
			)
			if test.authorize {
				req.Header.Set("Authorization", "Bearer "+sessionToken)
			}
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)

			require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String())
			require.Contains(t, rec.Body.String(), `"code":"INVALID_QUERY"`)
			require.Zero(t, upstreamCalls.Load())
		})
	}
}

// TestGatewayProxyDoesNotInheritControlTotalTimeout proves the P1
// client policy split: even when the underlying connection's
// control-plane client has a tiny total Timeout (50ms, simulating
// the production 30s cap), the gateway client must not inherit it —
// a 200ms streaming response must be fully proxied. Run time is
// bounded by the stream length, NOT a real 30s wait.
func TestGatewayProxyDoesNotInheritControlTotalTimeout(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Stream 4 chunks with 50ms gaps. Total stream time ≈200ms,
		// far exceeding the control client's 50ms total timeout.
		for i := range 4 {
			_, _ = w.Write([]byte("chunk-" + string(rune('1'+i))))
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			time.Sleep(50 * time.Millisecond)
		}
	}))
	defer upstream.Close()
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	// Set a 50ms total timeout on the control-plane client to prove
	// the gateway client does not inherit it. If it did, the 200ms
	// stream would fail with a context deadline exceeded.
	connections["server-a"].httpClient = &http.Client{
		Timeout:   50 * time.Millisecond,
		Transport: upstream.Client().Transport,
	}
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(http.MethodGet, "/api/v1/stream", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
	require.Equal(t, "chunk-1chunk-2chunk-3chunk-4", rec.Body.String(),
		"gateway must proxy the full stream despite the control client's total timeout")
}

// TestGatewayProxyNilTransportResponseHeaderTimeoutReturns502 proves
// the P1 fix for the production wire.New shape: control client has
// Transport=nil (so http.DefaultTransport is used implicitly). The
// gateway client must clone DefaultTransport and set a bounded
// ResponseHeaderTimeout so an upstream that accepts but never sends
// headers fails fast instead of hanging forever.
func TestGatewayProxyNilTransportResponseHeaderTimeoutReturns502(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Accept the request but never send headers. Simulates a
		// wedged upstream. The gateway's ResponseHeaderTimeout must
		// fire and target as a 502 to the client.
		<-r.Context().Done()
	}))
	defer upstream.Close()
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	// Mirror wire.New's httpClient shape: Timeout set, Transport=nil.
	connections["server-a"].httpClient = &http.Client{
		Timeout: 30 * time.Second,
	}
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(http.MethodGet, "/api/v1/never-responds", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()

	start := time.Now()
	router.ServeHTTP(rec, req)
	elapsed := time.Since(start)

	require.Equal(t, http.StatusBadGateway, rec.Code, rec.Body.String())
	require.Less(t, elapsed, 5*time.Second,
		"ResponseHeaderTimeout must bound the no-header case; got %v", elapsed)
}

// TestGatewayProxyClientDisconnectPropagatesToUpstream verifies that
// a REAL client socket disconnect propagates through the gateway to
// the upstream's request context. The prior version of this test
// used httptest.NewRequest + manual ctx.Cancel, which only proved the
// context tree wired correctly; this version drives a real HTTP
// listener + real http.Client so the disconnect goes through real
// net/conn state transitions.
//
// Evidence chain:
//   - upstream handler-entered channel fires (no sleep-based timing).
//   - client resp.Body.Close() with DisableKeepAlives tears the TCP
//     socket down immediately.
//   - the upstream produces data continuously so the gateway's
//     io.Copy keeps writing to c.Writer; the next Write after the
//     client close fails with EPIPE.
//   - io.Copy returns, proxy() defer-fires cancelRequest, the
//     standard *http.Transport tears down the upstream conn.
//   - upstream observes r.Context().Done() and exits.
//   - gatewayListener's ConnState hook eventually observes the
//     gateway-side connection transitioning to StateClosed.
func TestGatewayProxyClientDisconnectPropagatesToUpstream(t *testing.T) {
	upstreamEntered := make(chan struct{})
	upstreamContextDone := make(chan struct{})
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		close(upstreamEntered)
		w.WriteHeader(http.StatusOK)
		if f, ok := w.(http.Flusher); ok {
			f.Flush()
		}
		// Stream continuously to drive real-time propagation: the
		// gateway's io.Copy produces a Write on c.Writer per chunk.
		// The next Write after the client disconnect fails with
		// EPIPE, io.Copy returns, cancelRequest fires, the Transport
		// interrupts this Read, and r.Context() fires Done.
		ticker := time.NewTicker(20 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				if _, err := w.Write([]byte("chunk\n")); err != nil {
					close(upstreamContextDone)
					return
				}
				if f, ok := w.(http.Flusher); ok {
					f.Flush()
				}
			case <-r.Context().Done():
				close(upstreamContextDone)
				return
			}
		}
	}))
	defer upstream.Close()
	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	// Wrap the gin handler in a real http.Server so the client
	// request traverses a real listener. Track ConnState transitions
	// so the test can prove the gateway-side socket actually closed.
	// Use NewUnstartedServer so we can install ConnState BEFORE serve
	// starts; otherwise the serve goroutine races with the assignment.
	gatewayListener := httptest.NewUnstartedServer(router)
	var gatewayConnClosed atomic.Int32
	gatewayListener.Config.ConnState = func(_ net.Conn, state http.ConnState) {
		if state == http.StateClosed {
			gatewayConnClosed.Add(1)
		}
	}
	gatewayListener.Start()
	defer gatewayListener.Close()

	// DisableKeepAlives so resp.Body.Close() tears the TCP conn down
	// immediately instead of returning it to the client pool.
	disconnectClient := &http.Client{
		Transport: &http.Transport{DisableKeepAlives: true},
	}

	clientReq, err := http.NewRequest(http.MethodGet, gatewayListener.URL+"/api/v1/stream", nil)
	require.NoError(t, err)
	clientReq.Header.Set("Authorization", "Bearer "+sessionToken)

	resp, err := disconnectClient.Do(clientReq)
	require.NoError(t, err)
	require.NotNil(t, resp)

	// Wait until the upstream handler has actually entered — no sleep.
	select {
	case <-upstreamEntered:
	case <-time.After(3 * time.Second):
		_ = resp.Body.Close()
		t.Fatal("timed out waiting for upstream to enter handler")
	}

	// Read at least one chunk so we know the conn is fully
	// established end-to-end, then close the body. With
	// DisableKeepAlives=true the client transport tears the socket
	// down; the gateway's next Write fails with EPIPE; io.Copy
	// returns; cancelRequest fires; the Transport interrupts the
	// in-flight upstream Read; and the upstream observes
	// r.Context().Done().
	buf := make([]byte, 32)
	_, _ = resp.Body.Read(buf)
	_ = resp.Body.Close()

	select {
	case <-upstreamContextDone:
	case <-time.After(5 * time.Second):
		t.Fatal("upstream did not observe real client disconnect within 5s")
	}
	// After upstream observes disconnect, the gateway handler is
	// still unwinding: io.Copy returns its write error, proxy()'s
	// defer fires cancelRequest, the handler goroutine exits, and
	// finally the gateway listener's socket transitions StateClosed.
	require.Eventually(t, func() bool { return gatewayConnClosed.Load() > 0 },
		5*time.Second, 20*time.Millisecond,
		"gateway listener must observe its socket reaching StateClosed after client disconnect")
}

// TestGatewaySSEPreCommitRejectsNonFlusherLeaf verifies that when an
// SSE response (Content-Type: text/event-stream) is about to be
// committed but the leaf ResponseWriter does not implement
// http.Flusher, the gateway rejects the response pre-commit with a
// stable 5xx and never passes the upstream sentinel through.
//
// The wrapping writer structure is what reviewer3/codex5 mandated:
// two intermediate wrappers that EACH implement http.Flusher (so a
// single-level check would falsely pass) wrapping a non-Flusher leaf.
// leafFlusherOf must walk past both fakes, find no leaf Flusher, and
// the gateway must commit 5xx instead of starting a stream that
// would buffer forever.
//
// Red phase: the gateway currently has no SSE-specific path — it
// commits the upstream status verbatim and copies the body via
// io.Copy. The test fails because the response status is 200 (not
// 5xx) and the body contains the upstream sentinel.
// TestGatewaySSEPreCommitRejectsNonFlusherLeaf verifies that when the
// leaf writer does NOT implement http.Flusher (discovered via recursive
// Unwrap walk past two fake Flusher wrappers), the gateway MUST
// pre-commit reject with a stable 5xx BEFORE reading or forwarding any
// upstream response data.
//
// Assertions:
//   - status ≥ 500 (stable pre-commit rejection, not a panic-recovery 500)
//   - upstream sentinel does NOT leak into the response body
//   - upstream headers do NOT leak (X-Upstream-Leak absent)
//   - upstream connection is torn down within a bounded window
//     (r.Context().Done() fires, proving the gateway closed the body)
//
// Red phase: the gateway does not check the leaf Flusher before
// committing. It streams the upstream response (status=200, sentinel
// present) and leaves the upstream connection open.
// countingRoundTripper wraps an http.RoundTripper to count Read and
// Close calls on every response body that flows through. Used by the
// pre-commit non-Flusher test to PROVE the gateway did NOT consume
// upstream response body bytes before rejecting — `r.Context().Done`
// alone is insufficient evidence because http.Transport may have
// already pulled the body into its read buffer.
type countingRoundTripper struct {
	base     http.RoundTripper
	reads    atomic.Int64
	closes   atomic.Int64
	panicked atomic.Bool
}

func (rt *countingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := rt.base.RoundTrip(req)
	if err != nil {
		return nil, err
	}
	resp.Body = &countingReadCloser{
		Closer: resp.Body,
		onRead: func(n int, err error) {
			rt.reads.Add(int64(n))
			if err != nil && err != io.EOF {
				// non-EOF error: count as a real read that acked bytes
				// already accounted for by n.
			}
		},
		onClose: func() { rt.closes.Add(1) },
	}
	return resp, nil
}

type countingReadCloser struct {
	Closer  io.ReadCloser
	onRead  func(n int, err error)
	onClose func()
}

func (c *countingReadCloser) Read(p []byte) (int, error) {
	n, err := c.Closer.Read(p)
	if c.onRead != nil {
		c.onRead(n, err)
	}
	return n, err
}

func (c *countingReadCloser) Close() error {
	if c.onClose != nil {
		c.onClose()
	}
	return c.Closer.Close()
}

type gatewayResponseFunc func(*http.Request) (*http.Response, error)

func (f gatewayResponseFunc) RoundTrip(request *http.Request) (*http.Response, error) {
	return f(request)
}

type dataAndErrorReadCloser struct {
	done bool
}

func (r *dataAndErrorReadCloser) Read(p []byte) (int, error) {
	if r.done {
		return 0, io.EOF
	}
	r.done = true
	return copy(p, "data: trailing\n\n"), io.ErrUnexpectedEOF
}

func (*dataAndErrorReadCloser) Close() error { return nil }

func TestGatewaySSEFlushesTrailingDataBeforeUpstreamReadError(t *testing.T) {
	upstream := httptest.NewServer(http.NotFoundHandler())
	defer upstream.Close()
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{
		"server-a": upstream,
	})
	// The test server is used only for the stable server registry/origin.
	// Override transport behavior so the response body deterministically
	// returns (n>0, err) in one Read, which net/http transports rarely expose.
	connections["server-a"].gatewayClientMu.Lock()
	connections["server-a"].gatewayClient = &http.Client{
		Transport: gatewayResponseFunc(func(*http.Request) (*http.Response, error) {
			return &http.Response{
				StatusCode: http.StatusOK,
				Header: http.Header{
					"Content-Type": []string{"text/event-stream"},
				},
				Body: &dataAndErrorReadCloser{},
			}, nil
		}),
	}
	connections["server-a"].gatewayClientMu.Unlock()

	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(http.MethodGet, "/api/v1/events", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	leaf := &recordingFlusher{ResponseWriter: rec}
	router.ServeHTTP(leaf, req)

	require.Equal(t, "data: trailing\n\n", rec.Body.String(),
		"gateway must write n>0 bytes before returning the upstream read error")
	require.Equal(t, int32(1), leaf.flushes.Load(),
		"gateway must explicitly flush trailing n>0 bytes before returning the upstream read error")
}

type recordingFlushWriter struct {
	http.ResponseWriter
	flushes atomic.Int32
}

func (w *recordingFlushWriter) Flush() {
	w.flushes.Add(1)
	w.ResponseWriter.(http.Flusher).Flush()
}

func (w *recordingFlushWriter) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}

type gatedDownloadWriter struct {
	http.ResponseWriter
	firstWriteStarted chan struct{}
	releaseFirstWrite chan struct{}
	flushes           *atomic.Int32
	once              sync.Once
}

func (w *gatedDownloadWriter) Write(p []byte) (int, error) {
	w.once.Do(func() {
		close(w.firstWriteStarted)
		<-w.releaseFirstWrite
	})
	return w.ResponseWriter.Write(p)
}

func (w *gatedDownloadWriter) Flush() {
	if w.flushes != nil {
		w.flushes.Add(1)
	}
	w.ResponseWriter.(http.Flusher).Flush()
}

func TestGatewayOrdinaryDownloadPreservesBackpressureWithoutForcedFlush(t *testing.T) {
	firstChunk := bytes.Repeat([]byte("a"), 32*1024)
	secondChunk := bytes.Repeat([]byte("b"), 32*1024)
	payload := append(append([]byte(nil), firstChunk...), secondChunk...)
	secondChunkRead := make(chan struct{})
	upstreamBody := &gatedResponseBody{
		chunks:          [][]byte{firstChunk, secondChunk},
		secondChunkRead: secondChunkRead,
	}
	upstream := httptest.NewServer(http.NotFoundHandler())
	defer upstream.Close()
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	connections["server-a"].gatewayClientMu.Lock()
	connections["server-a"].gatewayClient = &http.Client{
		Transport: gatewayResponseFunc(func(*http.Request) (*http.Response, error) {
			return &http.Response{
				StatusCode: http.StatusOK,
				Header: http.Header{
					"Content-Type": []string{"application/octet-stream"},
				},
				Body: upstreamBody,
			}, nil
		}),
	}
	connections["server-a"].gatewayClientMu.Unlock()
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	firstWriteStarted := make(chan struct{})
	releaseFirstWrite := make(chan struct{})
	var downstreamFlushes atomic.Int32
	handler := http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		router.ServeHTTP(&gatedDownloadWriter{
			ResponseWriter:    writer,
			firstWriteStarted: firstWriteStarted,
			releaseFirstWrite: releaseFirstWrite,
			flushes:           &downstreamFlushes,
		}, request)
	})
	listener := httptest.NewServer(handler)
	defer listener.Close()

	req, err := http.NewRequest(http.MethodGet, listener.URL+"/api/v1/download", nil)
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	type result struct {
		resp *http.Response
		err  error
	}
	responseReady := make(chan result, 1)
	go func() {
		resp, requestErr := listener.Client().Do(req)
		responseReady <- result{resp: resp, err: requestErr}
	}()
	select {
	case <-firstWriteStarted:
	case <-time.After(time.Second):
		t.Fatal("gateway never reached the first downstream Write")
	}
	select {
	case <-secondChunkRead:
		t.Fatal("gateway read the second upstream chunk while the first downstream Write was blocked")
	default:
	}
	close(releaseFirstWrite)
	var resp *http.Response
	select {
	case result := <-responseReady:
		require.NoError(t, result.err)
		resp = result.resp
	case <-time.After(time.Second):
		t.Fatal("download did not resume after downstream became writable")
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, payload, body, "ordinary download lost or reordered data under backpressure")
	require.Zero(t, downstreamFlushes.Load(), "ordinary downloads must not use SSE-style forced flushes")
}

type gatedResponseBody struct {
	chunks          [][]byte
	next            int
	secondChunkRead chan struct{}
	closed          atomic.Bool
}

func (b *gatedResponseBody) Read(p []byte) (int, error) {
	if b.next >= len(b.chunks) {
		return 0, io.EOF
	}
	if b.next == 1 {
		close(b.secondChunkRead)
	}
	chunk := b.chunks[b.next]
	b.next++
	return copy(p, chunk), nil
}

func (b *gatedResponseBody) Close() error {
	b.closed.Store(true)
	return nil
}

func gatewayGenerationBridgeGoroutines() int {
	buffer := make([]byte, 1<<20)
	n := runtime.Stack(buffer, true)
	return strings.Count(
		string(buffer[:n]),
		"github.com/tidbcloud/AgentTeam/desktop/internal/api.(*desktopGateway).proxy.func1",
	)
}

func TestGatewayNormalEOFStopsGenerationBridge(t *testing.T) {
	upstreamEntered := make(chan struct{})
	releaseUpstream := make(chan struct{})
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		close(upstreamEntered)
		<-releaseUpstream
		_, _ = writer.Write([]byte("ok"))
	}))
	defer upstream.Close()
	router, _, registry := newGatewayTestRouterWithRegistry(
		t,
		map[string]*httptest.Server{"server-a": upstream},
	)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")
	baseline := gatewayGenerationBridgeGoroutines()

	done := make(chan *httptest.ResponseRecorder, 1)
	go func() {
		req := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
		req.Header.Set("Authorization", "Bearer "+sessionToken)
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)
		done <- rec
	}()
	select {
	case <-upstreamEntered:
	case <-time.After(time.Second):
		t.Fatal("long operation did not enter upstream")
	}
	require.Eventually(t, func() bool {
		return gatewayGenerationBridgeGoroutines() > baseline
	}, time.Second, 5*time.Millisecond,
		"long operation never installed a generation watcher goroutine")
	close(releaseUpstream)
	var rec *httptest.ResponseRecorder
	select {
	case rec = <-done:
	case <-time.After(time.Second):
		t.Fatal("gateway did not return after normal EOF")
	}
	require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
	require.Zero(t, registry.ActiveGatewayOps())
	require.Eventually(t, func() bool {
		return gatewayGenerationBridgeGoroutines() <= baseline
	}, time.Second, 5*time.Millisecond,
		"normal EOF/Release left generation watcher goroutines parked")
}

type cancelOnReadBody struct {
	entry <-chan *streamregistry.Entry
	cause error
	read  atomic.Bool
}

func (b *cancelOnReadBody) Read([]byte) (int, error) {
	if b.read.CompareAndSwap(false, true) {
		(<-b.entry).Cancel(b.cause)
	}
	return 0, errors.New("synthetic request body failure")
}

func (*cancelOnReadBody) Close() error { return nil }

type operationCancelCase struct {
	name          string
	cause         error
	wantStatus    int
	wantCode      string
	wantUnwritten bool
}

func operationCancelCases() []operationCancelCase {
	return []operationCancelCase{
		{
			name:       "generation",
			cause:      streamregistry.CancelError{Reason: streamregistry.ReasonGenerationChange},
			wantStatus: http.StatusConflict,
			wantCode:   "ACTIVE_SERVER_CHANGED",
		},
		{
			name:       "shutdown",
			cause:      streamregistry.CancelError{Reason: streamregistry.ReasonDaemonShutdown},
			wantStatus: http.StatusServiceUnavailable,
			wantCode:   "REGISTRY_CLOSED",
		},
		{
			name:          "client-disconnect",
			cause:         streamregistry.CancelError{Reason: streamregistry.ReasonClientDisconnect},
			wantUnwritten: true,
		},
		{
			name:          "deadline",
			cause:         context.DeadlineExceeded,
			wantUnwritten: true,
		},
	}
}

func requireOperationCancelResponse(t *testing.T, rec *httptest.ResponseRecorder, tc operationCancelCase) {
	t.Helper()
	if tc.wantUnwritten {
		require.NotContains(t, rec.Body.String(), "REMOTE_UNREACHABLE")
		require.NotContains(t, rec.Body.String(), "INVALID_REQUEST_BODY")
		require.NotContains(t, rec.Body.String(), "ACTIVE_SERVER_CHANGED")
		require.NotContains(t, rec.Body.String(), "REGISTRY_CLOSED")
		return
	}
	require.Equal(t, tc.wantStatus, rec.Code, rec.Body.String())
	require.Contains(t, rec.Body.String(), `"code":"`+tc.wantCode+`"`)
	require.NotContains(t, rec.Body.String(), "REMOTE_UNREACHABLE")
	require.NotContains(t, rec.Body.String(), "INVALID_REQUEST_BODY")
}

type countingResponseWriter struct {
	header      http.Header
	body        strings.Builder
	writeCalls  atomic.Int32
	headerCalls atomic.Int32
	status      atomic.Int32
}

func newBareGatewayHandler(
	connections types.ServerConnectionAccessor,
	sessions *uisession.Store,
	registry gatewayOperationRegistry,
) http.Handler {
	gateway := newDesktopGateway(connections, sessions, registry)
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		ctx, _ := gin.CreateTestContext(writer)
		ctx.Request = request
		// Mirror production handle(): WS preflight runs BEFORE any
		// auth/admission so a malformed or cross-origin WS intent returns
		// a 400/403 with zero side effects (no session lookup, no Registry
		// admission, no upstream dial). Without this, the bare handler would
		// diverge from production and WS rejection tests would not exercise
		// the same order.
		if isWebSocketIntent(ctx.Request) {
			if err := websocketPreflight(ctx.Request); err != nil {
				writeWebSocketPreflightError(ctx, err)
				return
			}
		}
		if sessionToken := strings.TrimPrefix(request.Header.Get("Authorization"), "Bearer "); sessionToken != "" {
			auth.AuthenticateSession(ctx, sessionToken)
		}
		gateway.proxy(ctx, gatewayQuery{})
	})
}

func newCountingResponseWriter() *countingResponseWriter {
	return &countingResponseWriter{header: make(http.Header)}
}

func (w *countingResponseWriter) Header() http.Header { return w.header }

func (w *countingResponseWriter) Write(p []byte) (int, error) {
	w.writeCalls.Add(1)
	return w.body.Write(p)
}

func (w *countingResponseWriter) WriteHeader(status int) {
	w.headerCalls.Add(1)
	w.status.Store(int32(status))
}

func requireNoGatewayResponse(t *testing.T, writer *countingResponseWriter) {
	t.Helper()
	require.Zero(t, writer.writeCalls.Load(), "client/deadline cancellation must not call Write")
	require.Zero(t, writer.headerCalls.Load(), "client/deadline cancellation must not call WriteHeader")
	require.Zero(t, writer.status.Load())
	require.Empty(t, writer.body.String())
	require.Empty(t, writer.header)
}

func serveGatewayCancelCase(
	t *testing.T,
	router http.Handler,
	sessionToken string,
	tc operationCancelCase,
) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	if !tc.wantUnwritten {
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)
		return rec
	}
	writer := newCountingResponseWriter()
	router.ServeHTTP(writer, req)
	requireNoGatewayResponse(t, writer)
	return nil
}

func TestGatewayBodyErrorPrefersTypedEntryCause(t *testing.T) {
	for _, tc := range operationCancelCases() {
		t.Run(tc.name, func(t *testing.T) {
			upstream := httptest.NewServer(http.NotFoundHandler())
			defer upstream.Close()
			realRegistry := streamregistry.New()
			entryReady := make(chan *streamregistry.Entry, 1)
			registry := &registerBarrierRegistry{
				delegate:   realRegistry,
				registered: entryReady,
			}
			router, _, store := newBareGatewayTestHandlerWithRegistry(
				t,
				map[string]*httptest.Server{"server-a": upstream},
				registry,
			)
			sessionToken := token.CreateSession()
			selectGatewayStore(t, store, sessionToken, "server-a")

			body := &cancelOnReadBody{entry: entryReady, cause: tc.cause}
			req := httptest.NewRequest(http.MethodPost, "/api/v1/upload", body)
			req.ContentLength = 1
			req.Header.Set("Authorization", "Bearer "+sessionToken)
			var (
				rec       *httptest.ResponseRecorder
				rawWriter *countingResponseWriter
				writer    http.ResponseWriter
			)
			if tc.wantUnwritten {
				rawWriter = newCountingResponseWriter()
				writer = rawWriter
			} else {
				rec = httptest.NewRecorder()
				writer = rec
			}
			done := make(chan struct{})
			go func() {
				defer close(done)
				router.ServeHTTP(writer, req)
			}()
			select {
			case <-done:
			case <-time.After(time.Second):
				t.Fatal("gateway did not return after body error")
			}
			if tc.wantUnwritten {
				requireNoGatewayResponse(t, rawWriter)
			} else {
				requireOperationCancelResponse(t, rec, tc)
			}
		})
	}
}

func TestGatewayTokenErrorPrefersTypedEntryCause(t *testing.T) {
	for _, tc := range operationCancelCases() {
		t.Run(tc.name, func(t *testing.T) {
			upstream := httptest.NewServer(http.NotFoundHandler())
			defer upstream.Close()
			realRegistry := streamregistry.New()
			entryReady := make(chan *streamregistry.Entry, 1)
			registry := &registerBarrierRegistry{
				delegate:   realRegistry,
				registered: entryReady,
			}
			router, connections, store := newBareGatewayTestHandlerWithRegistry(
				t,
				map[string]*httptest.Server{"server-a": upstream},
				registry,
			)
			connections["server-a"].tokenHook = func(context.Context) (string, error) {
				entry := <-entryReady
				entry.Cancel(tc.cause)
				return "", errors.New("synthetic token failure")
			}
			sessionToken := token.CreateSession()
			selectGatewayStore(t, store, sessionToken, "server-a")

			rec := serveGatewayCancelCase(t, router, sessionToken, tc)
			if rec != nil {
				requireOperationCancelResponse(t, rec, tc)
			}
		})
	}
}

func TestGatewayDialErrorPrefersTypedEntryCause(t *testing.T) {
	for _, tc := range operationCancelCases() {
		t.Run(tc.name, func(t *testing.T) {
			upstream := httptest.NewServer(http.NotFoundHandler())
			defer upstream.Close()
			realRegistry := streamregistry.New()
			entryReady := make(chan *streamregistry.Entry, 1)
			registry := &registerBarrierRegistry{
				delegate:   realRegistry,
				registered: entryReady,
			}
			router, connections, store := newBareGatewayTestHandlerWithRegistry(
				t,
				map[string]*httptest.Server{"server-a": upstream},
				registry,
			)
			connections["server-a"].gatewayClientMu.Lock()
			connections["server-a"].gatewayClient = &http.Client{
				Transport: gatewayResponseFunc(func(*http.Request) (*http.Response, error) {
					entry := <-entryReady
					entry.Cancel(tc.cause)
					return nil, errors.New("synthetic dial failure")
				}),
			}
			connections["server-a"].gatewayClientMu.Unlock()
			sessionToken := token.CreateSession()
			selectGatewayStore(t, store, sessionToken, "server-a")

			rec := serveGatewayCancelCase(t, router, sessionToken, tc)
			if rec != nil {
				requireOperationCancelResponse(t, rec, tc)
			}
		})
	}
}

func TestGatewayRefreshTokenErrorPrefersTypedEntryCause(t *testing.T) {
	for _, tc := range operationCancelCases() {
		t.Run(tc.name, func(t *testing.T) {
			upstream := httptest.NewServer(http.NotFoundHandler())
			defer upstream.Close()
			realRegistry := streamregistry.New()
			entryReady := make(chan *streamregistry.Entry, 1)
			registry := &registerBarrierRegistry{
				delegate:   realRegistry,
				registered: entryReady,
			}
			router, connections, store := newBareGatewayTestHandlerWithRegistry(
				t,
				map[string]*httptest.Server{"server-a": upstream},
				registry,
			)
			connections["server-a"].refreshHook = func(context.Context, string) (string, error) {
				current := <-entryReady
				current.Cancel(tc.cause)
				return "", errors.New("synthetic refresh failure")
			}
			connections["server-a"].gatewayClientMu.Lock()
			connections["server-a"].gatewayClient = &http.Client{
				Transport: gatewayResponseFunc(func(*http.Request) (*http.Response, error) {
					return &http.Response{
						StatusCode: http.StatusUnauthorized,
						Header:     make(http.Header),
						Body:       io.NopCloser(strings.NewReader("unauthorized")),
					}, nil
				}),
			}
			connections["server-a"].gatewayClientMu.Unlock()
			sessionToken := token.CreateSession()
			selectGatewayStore(t, store, sessionToken, "server-a")

			rec := serveGatewayCancelCase(t, router, sessionToken, tc)
			if rec != nil {
				requireOperationCancelResponse(t, rec, tc)
			}
		})
	}
}

func TestGatewayRetryDialErrorPrefersTypedEntryCause(t *testing.T) {
	for _, tc := range operationCancelCases() {
		t.Run(tc.name, func(t *testing.T) {
			upstream := httptest.NewServer(http.NotFoundHandler())
			defer upstream.Close()
			realRegistry := streamregistry.New()
			entryReady := make(chan *streamregistry.Entry, 1)
			registry := &registerBarrierRegistry{
				delegate:   realRegistry,
				registered: entryReady,
			}
			router, connections, store := newBareGatewayTestHandlerWithRegistry(
				t,
				map[string]*httptest.Server{"server-a": upstream},
				registry,
			)
			var calls atomic.Int32
			connections["server-a"].refreshToken = "fresh-user-token"
			connections["server-a"].gatewayClientMu.Lock()
			connections["server-a"].gatewayClient = &http.Client{
				Transport: gatewayResponseFunc(func(*http.Request) (*http.Response, error) {
					if calls.Add(1) == 1 {
						return &http.Response{
							StatusCode: http.StatusUnauthorized,
							Header:     make(http.Header),
							Body:       io.NopCloser(strings.NewReader("unauthorized")),
						}, nil
					}
					entry := <-entryReady
					entry.Cancel(tc.cause)
					return nil, errors.New("synthetic retry dial failure")
				}),
			}
			connections["server-a"].gatewayClientMu.Unlock()
			sessionToken := token.CreateSession()
			selectGatewayStore(t, store, sessionToken, "server-a")

			rec := serveGatewayCancelCase(t, router, sessionToken, tc)
			if rec != nil {
				requireOperationCancelResponse(t, rec, tc)
			}
			require.Equal(t, int32(2), calls.Load(), "test must reach retry doRequest")
		})
	}
}

func TestGatewayCopyErrorDoesNotOverwriteCommittedTypedCause(t *testing.T) {
	for _, tc := range operationCancelCases() {
		t.Run(tc.name, func(t *testing.T) {
			upstream := httptest.NewServer(http.NotFoundHandler())
			defer upstream.Close()
			realRegistry := streamregistry.New()
			entryReady := make(chan *streamregistry.Entry, 1)
			registry := &registerBarrierRegistry{
				delegate:   realRegistry,
				registered: entryReady,
			}
			router, connections, _ := newGatewayTestRouterWithStoreAndRegistry(
				t,
				map[string]*httptest.Server{"server-a": upstream},
				registry,
			)
			connections["server-a"].gatewayClientMu.Lock()
			connections["server-a"].gatewayClient = &http.Client{
				Transport: gatewayResponseFunc(func(*http.Request) (*http.Response, error) {
					entry := <-entryReady
					return &http.Response{
						StatusCode: http.StatusOK,
						Header:     make(http.Header),
						Body: &cancelOnReadResponseBody{
							entry: entry,
							cause: tc.cause,
						},
					}, nil
				}),
			}
			connections["server-a"].gatewayClientMu.Unlock()
			sessionToken := token.CreateSession()
			selectGatewayServer(t, router, sessionToken, "server-a")

			req := httptest.NewRequest(http.MethodGet, "/api/v1/download", nil)
			req.Header.Set("Authorization", "Bearer "+sessionToken)
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)
			require.Equal(t, http.StatusOK, rec.Code,
				"headers were committed before the copy error; typed cause must only terminate the stream")
			require.Empty(t, rec.Body.String())
			require.NotContains(t, rec.Body.String(), "REMOTE_UNREACHABLE")
			require.NotContains(t, rec.Body.String(), "ACTIVE_SERVER_CHANGED")
			require.NotContains(t, rec.Body.String(), "REGISTRY_CLOSED")
		})
	}
}

type cancelOnReadResponseBody struct {
	entry *streamregistry.Entry
	cause error
	once  sync.Once
}

func (b *cancelOnReadResponseBody) Read([]byte) (int, error) {
	b.once.Do(func() {
		b.entry.Cancel(b.cause)
	})
	return 0, errors.New("synthetic copy failure")
}

func (*cancelOnReadResponseBody) Close() error { return nil }

// countingReader wraps an io.Reader and counts bytes delivered via
// Read. Used by the B barrier test to prove the gateway consumed
// request body bytes after Register admitted the entry — the
// post-Register / pre-UserToken gap the B barrier must protect.
type countingReader struct {
	r      io.Reader
	onRead func(n int)
}

func (c *countingReader) Read(p []byte) (int, error) {
	n, err := c.r.Read(p)
	if n > 0 && c.onRead != nil {
		c.onRead(n)
	}
	return n, err
}

// TestGatewaySSEPreCommitRejectsNonFlusherLeaf verifies the gateway
// rejects a streaming response BEFORE committing a 200 to the client
// when the underlying leaf ResponseWriter does NOT implement a real
// http.Flusher. Two flusherWrappers hide a non-Flusher leaf; only a
// recursive unwrap can detect this.
//
// Pre-commit rejection contract:
//   - Status code: stable 5xx (gateway-owned; never leak upstream 200).
//   - Upstream response body: zero bytes consumed by the gateway
//     (counting RoundTripper proves Read count = 0). The previous
//     `r.Context().Done` evidence was insufficient because http.Transport
//     may have already pulled the body into its read buffer before the
//     upstream handler observed the context cancel.
//   - Upstream response body: Close count = 1 (gateway closed cleanly
//     exactly once; no leak, no double-close panic).
//   - Upstream sentinel bytes and X-Upstream-Leak header: NOT visible
//     to the client.
//   - Upstream connection: torn down (r.Context().Done observed).
//   - No panic from gateway, http.Server, or stubs.
func TestGatewaySSEPreCommitRejectsNonFlusherLeaf(t *testing.T) {
	const sentinel = "data:sse-precommit-sentinel\n\n"
	upstreamContextDone := make(chan struct{})
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
		w.Header().Set("X-Upstream-Leak", "should-not-appear")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(sentinel))
		if f, ok := w.(http.Flusher); ok {
			f.Flush()
		}
		select {
		case <-r.Context().Done():
			close(upstreamContextDone)
		case <-time.After(2 * time.Second):
		}
	}))
	defer upstream.Close()

	// countingRoundTripper sits between the gateway and the upstream
	// to count body Read/Close calls. We bypass BuildGatewayHTTPClient
	// (which rejects custom RoundTrippers by policy) and install a
	// gateway client directly: the same shape (Timeout=0, Jar=nil,
	// ErrUseLastResponse) but with our counting transport wrapping
	// upstream.Client().Transport.
	rt := &countingRoundTripper{
		base: upstream.Client().Transport,
	}
	router, connections := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	connections["server-a"].gatewayClientMu.Lock()
	connections["server-a"].gatewayClient = &http.Client{
		Timeout: 0,
		Jar:     nil,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
		Transport: rt,
	}
	connections["server-a"].gatewayClientMu.Unlock()
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	// Two fake-Flusher layers wrapping a non-Flusher leaf. The leaf
	// deliberately lacks Flush() AND Unwrap() so leafFlusherOf's
	// recursive walk terminates and returns (nil, false).
	wrap := func(real http.ResponseWriter) http.ResponseWriter {
		leaf := &struct{ http.ResponseWriter }{ResponseWriter: real}
		inner := &flusherWrapper{ResponseWriter: leaf}
		outer := &flusherWrapper{ResponseWriter: inner}
		return outer
	}
	wrappedRouter := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		router.ServeHTTP(wrap(w), r)
	})
	listener := httptest.NewServer(wrappedRouter)
	defer listener.Close()

	req, err := http.NewRequest(http.MethodGet, listener.URL+"/api/v1/sse", nil)
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()

	firstChunk := make([]byte, 256)
	n, _ := resp.Body.Read(firstChunk)
	bodySnippet := string(firstChunk[:n])

	require.Equalf(t, http.StatusBadGateway, resp.StatusCode,
		"expected stable 502 pre-commit rejection; got status=%d body=%q",
		resp.StatusCode, bodySnippet)
	require.Contains(t, bodySnippet, `"code":"SSE_NOT_SUPPORTED"`,
		"pre-commit rejection must use the stable SSE_NOT_SUPPORTED contract")
	require.NotContains(t, bodySnippet, sentinel,
		"sentinel leaked through pre-commit rejection")
	require.Empty(t, resp.Header.Get("X-Upstream-Leak"),
		"upstream header leaked through pre-commit rejection")

	// Upstream connection MUST be torn down: the gateway closes the
	// upstream response body after pre-commit, which fires
	// r.Context().Done() on the upstream side.
	select {
	case <-upstreamContextDone:
	case <-time.After(time.Second):
		t.Fatal("upstream context not cancelled after pre-commit rejection; gateway leaked the upstream connection")
	}

	// Stronger body-leak evidence: the gateway must NOT have read any
	// upstream response body bytes (Read count = 0) and must have
	// Closed the body exactly once. Red phase: gateway does
	// io.Copy(c.Writer, ...) which consumes the body — Read count > 0.
	if reads := rt.reads.Load(); reads != 0 {
		t.Fatalf("upstream response body was consumed by gateway: Read=%d bytes, want 0 (pre-commit must reject before any body read)",
			reads)
	}
	if closes := rt.closes.Load(); closes != 1 {
		t.Fatalf("upstream response body Close count = %d, want exactly 1 (clean single close, no leak, no double-close panic)",
			closes)
	}
	if rt.panicked.Load() {
		t.Fatal("panic observed in counting RoundTripper path")
	}
}

// TestGatewaySSEPerChunkViaChannelGate verifies per-chunk SSE delivery
// using a real client-side channel gate. The upstream writes chunk N,
// flushes, then blocks until the client confirms receipt by signalling
// `proceed`. Only then does the upstream write chunk N+1.
//
// This is strictly stronger than counting Flush() calls: if the gateway
// buffers chunks (no per-chunk leaf Flush), the data stays in the
// server's bufio buffer (net/http's bufferBeforeChunkedMode = 2048),
// the client's response-header read times out, and the test fails
// before any chunks can be exchanged.
//
// The listener deliberately uses the gateway router WITHOUT wrapping
// the writer in stubs. Wrapping via interface-embedding loses optional
// http.ResponseWriter interfaces (Hijacker, ReaderFrom, etc.) that
// net/http inspects; that perturbs the response framing — Go falls
// out of its small-body Content-Length optimization — and produces
// misleading failure modes (premature EOF after partial delivery
// instead of clean timeout).
//
// Both header-read and per-chunk body-read use bounded timeouts so the
// test fails fast in RED phase (gateway buffers, headers never arrive
// within 2s) and passes in GREEN phase (gateway leaf.Flush() per chunk,
// bytes arrive within milliseconds).
func TestGatewaySSEPerChunkViaChannelGate(t *testing.T) {
	const chunks = 3
	proceed := make(chan struct{})

	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		for i := 0; i < chunks; i++ {
			_, _ = w.Write([]byte("data: chunk\n\n"))
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			if i < chunks-1 {
				<-proceed
			}
		}
	}))
	defer upstream.Close()

	router, _ := newGatewayTestRouter(t, map[string]*httptest.Server{"server-a": upstream})
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	// No wrapWriter: the listener handler is the bare router so the
	// real http.ResponseWriter (which implements Flusher natively)
	// reaches the gateway. The channel gate is the sole failure
	// signal — leaf Flush counts would require a stub that loses
	// optional interfaces net/http relies on.
	listener := httptest.NewServer(router)
	defer listener.Close()

	// close(proceed) MUST run before listener.Close() and
	// upstream.Close(); otherwise the upstream handler stays parked
	// on <-proceed and httptest.Server.Close deadlocks for 5s then
	// leaks goroutines. Defer LIFO ordering: this defer is
	// registered LAST, so it runs FIRST.
	proceedClosed := false
	defer func() {
		if !proceedClosed {
			close(proceed)
		}
	}()

	req, err := http.NewRequest(http.MethodGet, listener.URL+"/api/v1/sse", nil)
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+sessionToken)

	// Phase 1: bounded header read. In RED phase, gateway's first
	// Write of 13 bytes stays in net/http's bufio buffer (< 2048
	// bytes, no Flush), so headers never reach the wire and
	// client.Do blocks forever. Cap it.
	type doResult struct {
		resp *http.Response
		err  error
	}
	doCh := make(chan doResult, 1)
	go func() {
		resp, err := listener.Client().Do(req)
		doCh <- doResult{resp, err}
	}()
	var resp *http.Response
	select {
	case r := <-doCh:
		require.NoError(t, r.err, "client.Do returned error")
		resp = r.resp
	case <-time.After(2 * time.Second):
		t.Fatal("client did not receive response headers within 2s; gateway not flushing SSE response (headers stuck in bufio)")
	}
	defer resp.Body.Close()

	// Phase 2: per-chunk channel gate using bufio.Reader. The reader
	// accumulates bytes until it can return a complete SSE event
	// (terminated by `\n\n`). Upstream writes chunk N (one event),
	// Flush, then blocks on <-proceed until we signal. Only after the
	// client observes the FULL event (including the trailing blank
	// line) do we release the upstream to write chunk N+1.
	//
	// This is strictly stronger than a raw byte Read:
	//   - Raw Read may return a partial event (e.g. "data: chu") when
	//     the gateway forwards bytes without flushing the terminator.
	//     A naive partial-read gate would falsely "see" data and signal
	//     proceed, letting the upstream write the next chunk on top of
	//     an unflushed terminator.
	//   - bufio.Reader's ReadString('\n') blocks until the newline
	//     arrives; a complete SSE event is two consecutive reads of
	//     '\n' (one ends the data line, one is the blank separator).
	//     If the gateway buffers the terminator, ReadString blocks,
	//     the bounded select times out, and the test fails.
	br := bufio.NewReader(resp.Body)
	for i := 0; i < chunks; i++ {
		type eventResult struct {
			data string
			err  error
		}
		ch := make(chan eventResult, 1)
		go func() {
			// ReadString('\n') returns one line at a time. For
			// "data: chunk\n\n" we expect:
			//   - line 1: "data: chunk\n"
			//   - line 2: "\n"  (blank separator)
			// Two consecutive reads complete one event.
			line1, err1 := br.ReadString('\n')
			if err1 != nil {
				ch <- eventResult{"", err1}
				return
			}
			line2, err2 := br.ReadString('\n')
			ch <- eventResult{line1 + line2, err2}
		}()
		select {
		case res := <-ch:
			if !strings.HasSuffix(res.data, "\n\n") {
				t.Fatalf("chunk %d: client received incomplete SSE event %q (err=%v status=%d); gateway not flushing the event terminator",
					i, res.data, res.err, resp.StatusCode)
			}
			// io.Reader may return (event, EOF) when the final chunk
			// completes the stream. Only empty counts as "no data".
			if res.data == "\n\n" {
				t.Fatalf("chunk %d: client received empty event (err=%v status=%d); gateway not flushing per SSE chunk",
					i, res.err, resp.StatusCode)
			}
			if res.err != nil && res.err != io.EOF && i < chunks-1 {
				t.Fatalf("chunk %d: stream ended prematurely after %q (err=%v); expected %d chunks total",
					i, res.data, res.err, chunks)
			}
		case <-time.After(time.Second):
			t.Fatalf("chunk %d: client did not receive a complete SSE event within 1s; gateway not flushing per-chunk terminator (status=%d)",
				i, resp.StatusCode)
		}
		if i < chunks-1 {
			select {
			case proceed <- struct{}{}:
			case <-time.After(time.Second):
				t.Fatalf("chunk %d: proceed signal blocked; upstream not consuming (deadlock)", i)
			}
		}
	}
}

// setupRealListenerGateway wraps the gateway router behind a REAL
// http.Server listener so ResponseController.SetReadDeadline can
// reach the actual TCP connection. Tests that exercise the
// read-deadline interrupt path MUST use a real listener —
// httptest.NewRecorder does not support SetReadDeadline and would
// force the implementation down a fake-hook path.
//
// wrapWriter, when non-nil, wraps each connection's ResponseWriter
// (e.g. to refuse SetReadDeadline for the non-capable path).
func setupRealListenerGateway(
	t *testing.T,
	upstreams map[string]*httptest.Server,
	wrapWriter func(http.ResponseWriter) http.ResponseWriter,
) (http.Handler, *streamregistry.Registry, *httptest.Server, string) {
	t.Helper()
	router, _, registry := newGatewayTestRouterWithRegistry(t, upstreams)
	sessionToken := token.CreateSession()
	for id := range upstreams {
		selectGatewayServer(t, router, sessionToken, id)
		break
	}
	handler := http.Handler(router)
	if wrapWriter != nil {
		wrapped := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			router.ServeHTTP(wrapWriter(w), r)
		})
		handler = wrapped
	}
	listener := httptest.NewServer(handler)
	t.Cleanup(listener.Close)
	return router, registry, listener, sessionToken
}

// startSlowUploadPipe sends a POST whose body is an io.Pipe that
// delivers exactly 1 byte then blocks. req.ContentLength is set so
// the gateway's prepareGatewayRequestBody enters the io.ReadAll
// buffering path — the code path where SetReadDeadline is required
// to interrupt a blocked read.
//
// Returns handlerDone (closes when the HTTP round-trip completes) and
// a release func that unblocks the pipe so the transport can unwind.
func startSlowUploadPipe(
	t *testing.T,
	listener *httptest.Server,
	sessionToken, path string,
) (<-chan struct{}, func()) {
	t.Helper()
	pr, pw := io.Pipe()
	releaseCh := make(chan struct{})
	go func() {
		_, _ = pw.Write([]byte("x"))
		<-releaseCh
		pw.Close()
	}()
	req, err := http.NewRequest(http.MethodPost, listener.URL+path, pr)
	require.NoError(t, err)
	req.ContentLength = 1024
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	req.Header.Set("Content-Type", "application/octet-stream")
	done := make(chan struct{})
	go func() {
		defer close(done)
		resp, err := listener.Client().Do(req)
		if err == nil {
			_, _ = io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
		}
	}()
	var once sync.Once
	return done, func() {
		once.Do(func() { close(releaseCh) })
	}
}

// waitForActiveOps waits until registry.ActiveGatewayOps() equals
// want or timeout expires. Returns false on timeout.
func waitForActiveOps(t *testing.T, registry *streamregistry.Registry, want int64, timeout time.Duration, msg string) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	ticker := time.NewTicker(5 * time.Millisecond)
	defer ticker.Stop()
	for {
		if registry.ActiveGatewayOps() == want {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("%s; activeOps=%d want=%d", msg, registry.ActiveGatewayOps(), want)
		}
		<-ticker.C
	}
}

// TestGatewaySlowUploadInterruptedByGenerationChangeReleasesEntry
// verifies that a generation change during a slow upload MUST release
// the gateway's registry entry within a bounded window via
// ResponseController.SetReadDeadline on the real TCP connection.
//
// The gateway's prepareGatewayRequestBody calls io.ReadAll on the
// request body; when the client is trickling bytes via io.Pipe,
// io.ReadAll parks on request.Body.Read. In the green phase, the
// generation change triggers entry.Cancel(generation-change), the
// bridge goroutine fires, SetReadDeadline(time.Now()) interrupts the
// read, io.ReadAll returns, and the handler unwinds → entry.Release().
//
// Red phase: no bridge exists; io.ReadAll stays parked; activeOps
// never returns to 0 within the 500ms bounded window.
func TestGatewaySlowUploadInterruptedByGenerationChangeReleasesEntry(t *testing.T) {
	upstreamA := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
	}))
	defer upstreamA.Close()
	upstreamB := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
	}))
	defer upstreamB.Close()

	router, registry, listener, sessionToken := setupRealListenerGateway(t,
		map[string]*httptest.Server{"server-a": upstreamA, "server-b": upstreamB},
		nil,
	)

	handlerDone, release := startSlowUploadPipe(t, listener, sessionToken, "/api/v1/upload")
	t.Cleanup(func() {
		release()
		select {
		case <-handlerDone:
		case <-time.After(3 * time.Second):
		}
	})

	waitForActiveOps(t, registry, 1, 2*time.Second, "entry was not registered")

	// Generation change: green phase calls entry.Cancel(generation-
	// change) which fires the bridge goroutine's SetReadDeadline.
	selectGatewayServer(t, router, sessionToken, "server-b")

	waitForActiveOps(t, registry, 0, 500*time.Millisecond,
		"entry not released after generation change; SetReadDeadline bridge needed")
}

// TestGatewaySlowUploadInterruptedByCancelAllReleasesEntry verifies
// that registry.CancelAll — the daemon-shutdown drain path — also
// releases the entry via SetReadDeadline on the real connection.
//
// Red phase: CancelAll cancels entry.Context(), but without a bridge
// goroutine watching entry.Context().Done(), the io.ReadAll stays
// parked and activeOps never returns to 0.
func TestGatewaySlowUploadInterruptedByCancelAllReleasesEntry(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
	}))
	defer upstream.Close()

	_, registry, listener, sessionToken := setupRealListenerGateway(t,
		map[string]*httptest.Server{"server-a": upstream},
		nil,
	)

	handlerDone, release := startSlowUploadPipe(t, listener, sessionToken, "/api/v1/upload")
	t.Cleanup(func() {
		release()
		select {
		case <-handlerDone:
		case <-time.After(3 * time.Second):
		}
	})

	waitForActiveOps(t, registry, 1, 2*time.Second, "entry was not registered")

	// CancelAll is the daemon-shutdown drain. In green, the bridge
	// goroutine reacts to entry.Context().Done() and fires
	// SetReadDeadline to interrupt the blocked io.ReadAll.
	_ = registry.CancelAll()

	waitForActiveOps(t, registry, 0, 500*time.Millisecond,
		"entry not released after CancelAll; SetReadDeadline bridge needed")
}

// TestGatewayNonCapableWriterCancelAllWaitTimesOut verifies the
// ErrNotSupported path: when the underlying writer does NOT support
// SetReadDeadline, the bridge cannot force-interrupt the blocked
// io.ReadAll. CancelAll returns its wait function immediately, but
// calling wait(ctx) with a bounded ctx MUST time out because the
// entry cannot be released.
//
// Red phase: no bridge exists at all, so SetReadDeadline is never
// attempted. The test still verifies that CancelAll returns a wait
// function and that wait(ctx) times out — the contract the daemon
// shutdown relies on to stay Preparing instead of deadlocking.
func TestGatewayNonCapableWriterCancelAllWaitTimesOut(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
	}))
	defer upstream.Close()

	// Wrap every connection's writer with deadlineRefusingWriter so
	// ResponseController.SetReadDeadline returns ErrNotSupported.
	_, registry, listener, sessionToken := setupRealListenerGateway(t,
		map[string]*httptest.Server{"server-a": upstream},
		func(real http.ResponseWriter) http.ResponseWriter {
			return &deadlineRefusingWriter{ResponseWriter: real}
		},
	)

	handlerDone, release := startSlowUploadPipe(t, listener, sessionToken, "/api/v1/upload")
	t.Cleanup(func() {
		release()
		select {
		case <-handlerDone:
		case <-time.After(3 * time.Second):
		}
	})

	waitForActiveOps(t, registry, 1, 2*time.Second, "entry was not registered")

	wait := registry.CancelAll()

	// CancelAll returns immediately; the returned wait function blocks
	// on entry release. In both red (no bridge) and green (bridge but
	// SetReadDeadline returns ErrNotSupported), the entry cannot be
	// released, so wait(ctx) MUST time out.
	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
	defer cancel()
	err := wait(ctx)
	if err == nil {
		t.Fatal("wait returned nil; expected context deadline exceeded when writer cannot honor SetReadDeadline")
	}
}

// ---------------------------------------------------------------------------
// A/B/C generation barriers
//
// Reviewer3/planner3 requirement: the gateway must defend THREE
// generation barriers around `registry.Register`, not just the existing
// pre-Watch check. Each barrier must be exercised against the REAL
// gateway+uisession+registry integration — not via "推定覆盖" through
// adjacent paths.
//
//   A (pre-Watch):       generation change before WatchGeneration is
//                        consulted → request rejected with 409, no entry
//                        admitted (activeOps stays 0).
//   B (post-Register):   generation change after Register succeeds but
//                        BEFORE the gateway consumes the request body,
//                        dials UserToken, or fires doRequest → 409 with
//                        ZERO side effects (token calls=0, body Read=0,
//                        upstream dial=0). Requires a post-Register
//                        recheck that does NOT exist today.
//   C (post-Register):   generation change AFTER Register and AFTER the
//                        upstream is in-flight → the entry MUST receive
//                        a typed ReasonGenerationChange cancel (not just
//                        requestContext cancel) and activeOps MUST
//                        return to 0. Requires the gateway to call
//                        registry.CancelSessionGeneration (or
//                        entry.Cancel) on generation change, which the
//                        current goroutine-only requestContext cancel
//                        does NOT do.
//
// Red phase: B and C fail. A passes (it tests the existing barrier).
// ---------------------------------------------------------------------------

// TestGatewayPreWatchBarrierRejectsWithoutAdmitting is the A barrier
// test. The uisession.Store is driven directly into a server-transition
// state, which makes the gateway's `WatchGeneration(sessionID, gen)`
// return ok=false BEFORE any Register call. The request MUST return
// 409 ACTIVE_SERVER_CHANGED and the registry MUST NOT admit any entry.
//
// Determinism: server transition is a stable, lock-protected state —
// we don't race the Get-vs-Watch window. We simply put the store into
// a state where WatchGeneration ok=false unconditionally for any
// generation, send the request, and assert the gateway's response.
func TestGatewayPreWatchBarrierRejectsWithoutAdmitting(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
	defer upstream.Close()

	router, _, registry, store := newGatewayTestRouterWithStore(t,
		map[string]*httptest.Server{"server-a": upstream},
	)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	// BeginServerTransition fences the server: WatchGeneration returns
	// ok=false for any generation while the transition is open. This is
	// the deterministic equivalent of "generation changed between
	// sessions.Get and sessions.WatchGeneration".
	endTransition := store.BeginServerTransition("server-a")
	defer endTransition()

	req := httptest.NewRequest(
		http.MethodPost,
		"/api/v1/example",
		strings.NewReader(`{"ok":true}`),
	)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equalf(t, http.StatusConflict, rec.Code,
		"A barrier must return 409 ACTIVE_SERVER_CHANGED during server transition; got status=%d body=%q",
		rec.Code, rec.Body.String())
	require.Contains(t, rec.Body.String(), "ACTIVE_SERVER_CHANGED",
		"A barrier response body must name the ACTIVE_SERVER_CHANGED code; got %q", rec.Body.String())

	// Critical A-barrier property: no entry was admitted. If the gateway
	// had called Register before consulting WatchGeneration, activeOps
	// would be ≥ 1 and the entry would leak when the handler returns.
	require.Equal(t, int64(0), registry.ActiveGatewayOps(),
		"A barrier must reject BEFORE Register; activeOps must stay 0")
}

// TestGatewayPostRegisterBarrierHasNoSideEffectOnGenerationChange is
// the B barrier test. A wrapper parks immediately BEFORE delegating to
// the real Registry.Register. While parked, the test advances the real
// uisession generation, then releases Register. The old-generation entry
// is admitted after the watch already fired, so only the mandatory
// post-register MatchesGeneration recheck can stop the handler before
// request-body/token/dial side effects.
func TestGatewayPostRegisterBarrierHasNoSideEffectOnGenerationChange(t *testing.T) {
	var upstreamCalls atomic.Int32
	upstreamA := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		upstreamCalls.Add(1)
		_, _ = io.Copy(io.Discard, r.Body)
	}))
	defer upstreamA.Close()
	upstreamB := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
	}))
	defer upstreamB.Close()

	realRegistry := streamregistry.New()
	beforeDelegate := make(chan struct{})
	releaseBefore := make(chan struct{})
	barrier := &registerBarrierRegistry{
		delegate:       realRegistry,
		beforeDelegate: beforeDelegate,
		releaseBefore:  releaseBefore,
	}
	router, connections, store := newGatewayTestRouterWithStoreAndRegistry(
		t,
		map[string]*httptest.Server{"server-a": upstreamA, "server-b": upstreamB},
		barrier,
	)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	var requestBytesRead atomic.Int64
	body := &countingReader{
		r: strings.NewReader("side-effect-must-not-run"),
		onRead: func(n int) {
			requestBytesRead.Add(int64(n))
		},
	}
	req := httptest.NewRequest(http.MethodPost, "/api/v1/upload", io.NopCloser(body))
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	req.Header.Set("Content-Type", "application/octet-stream")
	rec := httptest.NewRecorder()
	done := make(chan struct{})
	go func() {
		defer close(done)
		router.ServeHTTP(rec, req)
	}()

	select {
	case <-beforeDelegate:
	case <-time.After(2 * time.Second):
		t.Fatal("B barrier: gateway did not park before delegate Register")
	}
	current, ok := store.Get(sessionToken)
	require.True(t, ok)
	_, err := store.Select(sessionToken, current.ExpiresAt, "server-b")
	require.NoError(t, err)
	close(releaseBefore)
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("B barrier: handler did not return")
	}

	require.Equal(t, http.StatusConflict, rec.Code, rec.Body.String())
	require.Contains(t, rec.Body.String(), "ACTIVE_SERVER_CHANGED")
	require.Zero(t, requestBytesRead.Load(), "B barrier must reject before reading request body")
	require.Zero(t, connections["server-a"].tokenCalls.Load(), "B barrier must reject before UserToken")
	require.Zero(t, upstreamCalls.Load(), "B barrier must reject before upstream dial")
	require.Zero(t, realRegistry.ActiveGatewayOps(), "B barrier entry must release")
}

// TestGatewayPostRegisterBarrierPropagatesTypedCancelToEntry is the C
// barrier test. A wrapper delegates to the real Registry.Register and
// then parks BEFORE returning the admitted Entry to the gateway. The
// test advances generation in that exact interval, releases Register,
// and verifies the handler's own generation watcher cancels this Entry
// with ReasonGenerationChange. No registry-wide scan or diagnostic API
// is used.
func TestGatewayPostRegisterBarrierPropagatesTypedCancelToEntry(t *testing.T) {
	upstreamA := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
	defer upstreamA.Close()
	upstreamB := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
	defer upstreamB.Close()

	realRegistry := streamregistry.New()
	afterDelegate := make(chan struct{})
	releaseAfter := make(chan struct{})
	registered := make(chan *streamregistry.Entry, 1)
	barrier := &registerBarrierRegistry{
		delegate:      realRegistry,
		afterDelegate: afterDelegate,
		releaseAfter:  releaseAfter,
		registered:    registered,
	}
	router, _, store := newGatewayTestRouterWithStoreAndRegistry(
		t,
		map[string]*httptest.Server{"server-a": upstreamA, "server-b": upstreamB},
		barrier,
	)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	req := httptest.NewRequest(http.MethodGet, "/api/v1/example", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	done := make(chan struct{})
	go func() {
		defer close(done)
		router.ServeHTTP(rec, req)
	}()

	select {
	case <-afterDelegate:
	case <-time.After(2 * time.Second):
		t.Fatal("C barrier: gateway did not park after delegate Register")
	}
	var entry *streamregistry.Entry
	select {
	case entry = <-registered:
	case <-time.After(time.Second):
		t.Fatal("C barrier: admitted entry not observed")
	}
	require.Nil(t, context.Cause(entry.Context()))
	current, ok := store.Get(sessionToken)
	require.True(t, ok)
	_, err := store.Select(sessionToken, current.ExpiresAt, "server-b")
	require.NoError(t, err)
	close(releaseAfter)

	require.Eventually(t, func() bool {
		return errors.Is(context.Cause(entry.Context()), streamregistry.ReasonGenerationChange)
	}, time.Second, 5*time.Millisecond,
		"C barrier: current entry must receive typed generation-change cause")
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("C barrier: handler did not return")
	}
	require.Equal(t, http.StatusConflict, rec.Code, rec.Body.String())
	require.Zero(t, realRegistry.ActiveGatewayOps(), "C barrier entry must release")
}

// ---------------------------------------------------------------------------
// Daemon double-drain runtime contract
//
// Reviewer3 requirement: daemon shutdown must execute TWO sequential
// drains with bounded budgets, and the admission MUST be closed as
// soon as the first drain starts:
//
//   step 1: registry.CancelAll() — synchronous, fires typed
//           ReasonDaemonShutdown on every active entry. Admission
//           closes immediately: subsequent Register calls return
//           ErrRegistryClosed.
//   step 2: waitFn(ctx1) — bounded by ctx1. Returns nil when every
//           entry has been Released (i.e. every handler observed the
//           typed cancel and unwound). Returns ctx1.Err() if any
//           handler stays parked past ctx1's deadline.
//   step 3: http.Server.Shutdown(ctx2) — only safe to call after
//           waitFn has returned nil, because there are no active
//           handlers left to block Shutdown.
//
// State transitions:
//   - waitFn == nil AND http.Server.Shutdown == nil → Offline
//   - either timed out → stay Preparing (do NOT proceed to Offline)
//
// The full Daemon struct has too many filesystem/config dependencies
// to construct in a unit test, so this test exercises the contract at
// the gateway+listener level. The daemon code path itself is a thin
// wrapper that calls these primitives in the same order.
// ---------------------------------------------------------------------------

// TestGatewayDaemonDoubleDrainClosesAdmissionAndBoundsHTTPShutdown
// drives the full daemon shutdown sequence against a real listener
// with a parked slow upload:
//
//  1. CancelAll — admission closes immediately; a NEW gateway request
//     returns 503 REGISTRY_CLOSED.
//  2. waitFn(boundedCtx) — must return nil within budget. The
//     SetReadDeadline bridge (green phase) interrupts the parked
//     body Read; the handler unwinds; entry.Release fires; the
//     drained channel closes; waitFn returns nil.
//  3. listener.Shutdown(boundedCtx) — must return promptly with no
//     active handlers left.
//
// State transition contract: BOTH drains succeed → daemon reaches
// Offline. EITHER drain times out → daemon stays Preparing.
//
// RED phase: the bridge is not yet wired, so waitFn CANNOT drain the
// parked entry; the test fails at step 2 (waitFn returns
// DeadlineExceeded instead of nil). GREEN phase wires the bridge and
// the test passes end-to-end.
func TestGatewayDaemonDoubleDrainClosesAdmissionAndBoundsHTTPShutdown(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
	}))
	defer upstream.Close()

	router, _, registry, _ := newGatewayTestRouterWithStore(t,
		map[string]*httptest.Server{"server-a": upstream},
	)
	sessionToken := token.CreateSession()
	selectGatewayServer(t, router, sessionToken, "server-a")

	listener := httptest.NewServer(router)
	t.Cleanup(listener.Close)

	// Park a slow upload so the gateway handler is stuck on
	// io.ReadAll(request.Body).Registry has one active entry.
	handlerDone, release := startSlowUploadPipe(t, listener, sessionToken, "/api/v1/upload")
	t.Cleanup(func() {
		release()
		select {
		case <-handlerDone:
		case <-time.After(3 * time.Second):
		}
	})
	waitForActiveOps(t, registry, 1, 2*time.Second, "daemon drain: entry was not registered")

	// STEP 1: CancelAll — synchronous, fires typed cancel. The wait
	// function is the daemon's drain handle.
	wait := registry.CancelAll()

	// Admission is NOW closed: a new gateway request returns 503.
	newReq := httptest.NewRequest(http.MethodGet, "/api/v1/example", nil)
	newReq.Header.Set("Authorization", "Bearer "+sessionToken)
	newRec := httptest.NewRecorder()
	router.ServeHTTP(newRec, newReq)
	require.Equalf(t, http.StatusServiceUnavailable, newRec.Code,
		"closed registry must yield 503 REGISTRY_CLOSED; got status=%d body=%q",
		newRec.Code, newRec.Body.String())
	require.Contains(t, newRec.Body.String(), "REGISTRY_CLOSED",
		"503 response body must name the REGISTRY_CLOSED code; got %q", newRec.Body.String())
	// The read-deadline bridge may already have interrupted and released the
	// original entry by this point. Both activeOps=1 (wait still converging)
	// and activeOps=0 (fast-path drain already completed) are valid. The
	// bounded wait below is the synchronization point that proves the original
	// admitted operation has fully released.

	// STEP 2: bounded wait. GREEN phase: the SetReadDeadline bridge
	// interrupts the parked body Read; the handler unwinds; the entry
	// is released; the drained channel closes; waitFn returns nil
	// well before the 500ms budget.
	//
	// RED phase: no bridge, the entry stays parked, waitFn returns
	// DeadlineExceeded and this assertion fails.
	waitCtx, waitCancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer waitCancel()
	require.NoError(t, wait(waitCtx),
		"daemon drain wait must succeed within budget; SetReadDeadline bridge must interrupt parked body Read")

	// STEP 3: http.Server.Shutdown must return promptly. The green
	// waitFn drained every active handler, so Shutdown has nothing
	// to wait on.
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
	defer shutdownCancel()
	require.NoError(t, listener.Config.Shutdown(shutdownCtx),
		"http.Server.Shutdown must return promptly after registry drain; active handlers should be gone")
}