iop/apps/edge/internal/transport/integration_test.go

package transport_test

import (
	"context"
	"fmt"
	"net"
	"sync"
	"testing"
	"time"

	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"git.toki-labs.com/toki/proto-socket/go/packets"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	edgeservice "iop/apps/edge/internal/service"
	"iop/apps/edge/internal/transport"
	"iop/packages/go/config"
	eventpkg "iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

func getFreePort(t *testing.T) string {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	addr := l.Addr().String()
	l.Close()
	return addr
}

// handshakeAttemptTimeout / handshakeMaxAttempts bound one register/ready
// request and how many fresh connections the resilient helpers try. The repeated
// -race suite runs three race-instrumented binaries in parallel and occasionally
// stalls a fresh localhost connection's handshake for seconds while the machine
// is otherwise idle (a scheduling/netpoll artifact, not a product hang). Rather
// than inflate a single fixed request ceiling — which only makes a stalled
// attempt fail slower — the helpers abandon a stalled attempt quickly and retry
// on a fresh connection, which succeeds. Production heartbeat/timeout constants
// are untouched; these bound only the tests' own handshake.
const (
	handshakeAttemptTimeout = 5 * time.Second
	handshakeMaxAttempts    = 4
)

func nodeParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.RegisterResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.RegisterResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeReadyResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeReadyResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

func tryDialNode(ctx context.Context, listenAddr string) (*toki.TcpClient, error) {
	host, portStr, _ := net.SplitHostPort(listenAddr)
	port := 0
	fmt.Sscanf(portStr, "%d", &port)
	return toki.DialTcp(ctx, host, port, 30, 10, nodeParserMap())
}

func dialNode(t *testing.T, ctx context.Context, listenAddr string) *toki.TcpClient {
	t.Helper()
	client, err := tryDialNode(ctx, listenAddr)
	if err != nil {
		t.Fatalf("dial edge server: %v", err)
	}
	return client
}

// resilientRegister performs the dial+register handshake, retrying the whole
// thing on a fresh connection when an attempt errors (a transient connection
// stall under the repeated -race suite). A completed round-trip — accepted OR
// rejected — is returned as-is; only connection/transport errors are retried, so
// a genuine rejection (e.g. a duplicate registration) is never masked. The
// returned client is the one that carried the successful round-trip.
func resilientRegister(t *testing.T, ctx context.Context, listenAddr, token string) (*toki.TcpClient, *iop.RegisterResponse) {
	t.Helper()
	var lastErr error
	for attempt := 1; attempt <= handshakeMaxAttempts; attempt++ {
		client, err := tryDialNode(ctx, listenAddr)
		if err != nil {
			lastErr = err
			continue
		}
		resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator, &iop.RegisterRequest{Token: token}, handshakeAttemptTimeout)
		if err == nil {
			return client, resp
		}
		lastErr = err
		client.Close()
	}
	t.Fatalf("register handshake failed after %d attempts: %v", handshakeMaxAttempts, lastErr)
	return nil, nil
}

// signalReady drives the node→edge NodeReadyRequest handshake and asserts the
// edge marked the connection dispatch-ready. Registration only claims ownership
// and delivers config; the connect lifecycle and connected event fire on this
// signal, so tests that assert connectivity must call it after an accepted
// registration. The ready request is retried on the same (already-registered)
// connection when it stalls transiently; a duplicate ready is idempotent server
// side, so a retry is safe.
func signalReady(t *testing.T, client *toki.TcpClient, nodeID string) {
	t.Helper()
	var lastErr error
	for attempt := 1; attempt <= handshakeMaxAttempts; attempt++ {
		resp, err := toki.SendRequestTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
			&client.Communicator, &iop.NodeReadyRequest{NodeId: nodeID}, handshakeAttemptTimeout)
		if err != nil {
			lastErr = err
			if !client.IsAlive() {
				break
			}
			continue
		}
		if !resp.GetReady() {
			t.Fatalf("expected ready ack, got reason %q", resp.GetReason())
		}
		return
	}
	t.Fatalf("ready request failed after %d attempts: %v", handshakeMaxAttempts, lastErr)
}

// signalReadyResp is the raw single-shot ready handshake for tests that assert
// the response (duplicate/stale ready) rather than a successful ack.
func signalReadyResp(t *testing.T, client *toki.TcpClient, nodeID string) *iop.NodeReadyResponse {
	t.Helper()
	resp, err := toki.SendRequestTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
		&client.Communicator, &iop.NodeReadyRequest{NodeId: nodeID}, handshakeAttemptTimeout)
	if err != nil {
		t.Fatalf("ready request: %v", err)
	}
	return resp
}

func waitForRegistryEntry(ctx context.Context, registry *edgenode.Registry, nodeID string) (*edgenode.NodeEntry, bool) {
	ticker := time.NewTicker(25 * time.Millisecond)
	defer ticker.Stop()

	for {
		if entry, ok := registry.Get(nodeID); ok {
			return entry, true
		}

		select {
		case <-ctx.Done():
			return nil, false
		case <-ticker.C:
		}
	}
}

func waitForRegistryAbsent(ctx context.Context, registry *edgenode.Registry, nodeID string) bool {
	ticker := time.NewTicker(25 * time.Millisecond)
	defer ticker.Stop()

	for {
		if _, ok := registry.Get(nodeID); !ok {
			return true
		}

		select {
		case <-ctx.Done():
			return false
		case <-ticker.C:
		}
	}
}

func TestEdgeServerGenericRegistrationKind(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{
			ID:        "node-test-01",
			Alias:     "generic-node",
			Token:     "node-token",
			AgentKind: config.AgentKindGenericNode,
		},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	client := dialNode(t, ctx, listenAddr)
	defer client.Close()

	resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client.Communicator,
		&iop.RegisterRequest{Token: "node-token"},
		45*time.Second,
	)
	if err != nil {
		t.Fatalf("register request: %v", err)
	}
	if !resp.GetAccepted() {
		t.Fatalf("expected accepted register response, got reason %q", resp.GetReason())
	}

	entry, ok := waitForRegistryEntry(ctx, registry, "node-test-01")
	if !ok {
		t.Fatal("node was not registered in edge registry within timeout")
	}
	if entry.AgentKind != config.AgentKindGenericNode {
		t.Fatalf("agent kind: got %q want %q", entry.AgentKind, config.AgentKindGenericNode)
	}
	if entry.LifecycleState != edgenode.LifecycleConnected {
		t.Fatalf("lifecycle: got %q want %q", entry.LifecycleState, edgenode.LifecycleConnected)
	}
}

func TestEdgeServerDuplicateRegistrationReason(t *testing.T) {
	cases := []struct {
		name       string
		kind       string
		wantReason string
	}{
		{"generic", config.AgentKindGenericNode, "node already connected"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
			defer cancel()

			logger := zap.NewNop()
			listenAddr := getFreePort(t)
			registry := edgenode.NewRegistry()
			nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
				{ID: "dup-01", Alias: "dup", Token: "dup-token", AgentKind: tc.kind},
			})
			if err != nil {
				t.Fatalf("load node store: %v", err)
			}

			server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
			if err != nil {
				t.Fatalf("new server: %v", err)
			}
			if err := server.Start(ctx); err != nil {
				t.Fatalf("start server: %v", err)
			}
			defer server.Stop()

			first := dialNode(t, ctx, listenAddr)
			defer first.Close()
			resp1, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
				&first.Communicator, &iop.RegisterRequest{Token: "dup-token"}, 45*time.Second)
			if err != nil {
				t.Fatalf("first register: %v", err)
			}
			if !resp1.GetAccepted() {
				t.Fatalf("expected first registration accepted, got reason %q", resp1.GetReason())
			}
			if _, ok := waitForRegistryEntry(ctx, registry, "dup-01"); !ok {
				t.Fatal("first registration not visible in registry within timeout")
			}

			second := dialNode(t, ctx, listenAddr)
			defer second.Close()
			resp2, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
				&second.Communicator, &iop.RegisterRequest{Token: "dup-token"}, 45*time.Second)
			if err != nil {
				t.Fatalf("second register: %v", err)
			}
			if resp2.GetAccepted() {
				t.Fatal("expected second registration to be rejected")
			}
			if resp2.GetReason() != tc.wantReason {
				t.Fatalf("duplicate reason: got %q want %q", resp2.GetReason(), tc.wantReason)
			}
		})
	}
}

func TestEdgeServerReconnectAfterUnregisterAccepted(t *testing.T) {
	// Widened from 5s/2s so the race detector's instrumentation slowdown, observed
	// to trip a 2s request timeout once under the full -race suite, cannot flake
	// this deterministically-passing reconnect handshake. Production heartbeat
	// constants are untouched; only the test's own deadlines move.
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: "reconnect-01", Alias: "reconnect", Token: "reconnect-token", AgentKind: config.AgentKindGenericNode},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	first := dialNode(t, ctx, listenAddr)
	resp1, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&first.Communicator, &iop.RegisterRequest{Token: "reconnect-token"}, 45*time.Second)
	if err != nil {
		t.Fatalf("first register: %v", err)
	}
	if !resp1.GetAccepted() {
		t.Fatalf("expected first registration accepted, got reason %q", resp1.GetReason())
	}
	if _, ok := waitForRegistryEntry(ctx, registry, "reconnect-01"); !ok {
		t.Fatal("first registration not visible in registry within timeout")
	}

	if err := first.Close(); err != nil {
		t.Fatalf("close first client: %v", err)
	}
	if ok := waitForRegistryAbsent(ctx, registry, "reconnect-01"); !ok {
		t.Fatal("registry still contains node after disconnect")
	}

	second := dialNode(t, ctx, listenAddr)
	defer second.Close()
	resp2, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&second.Communicator, &iop.RegisterRequest{Token: "reconnect-token"}, 45*time.Second)
	if err != nil {
		t.Fatalf("second register: %v", err)
	}
	if !resp2.GetAccepted() {
		t.Fatalf("expected reconnect after unregister accepted, got reason %q", resp2.GetReason())
	}
	if _, ok := waitForRegistryEntry(ctx, registry, "reconnect-01"); !ok {
		t.Fatal("reconnect registration not visible in registry within timeout")
	}
}

func TestEdgeServerIntegration(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{
			ID:    "node-test-01",
			Alias: "test-node",
			Token: "test-token",
			Adapters: config.AdaptersConf{
				Ollama: config.OllamaConf{Enabled: true, BaseURL: "http://localhost:11434"},
				Mock:   config.MockConf{Enabled: true},
			},
			Runtime: config.RuntimeConf{Concurrency: 2},
		},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	// 1. Edge 서버 구동
	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("failed to create server: %v", err)
	}
	nodeEventCh := make(chan *iop.EdgeNodeEvent, 4)
	server.SetNodeEventHandler(func(event *iop.EdgeNodeEvent) {
		nodeEventCh <- event
	})
	if err := server.Start(ctx); err != nil {
		t.Fatalf("failed to start server: %v", err)
	}
	defer server.Stop()

	// 2. Mock Node 클라이언트 접속
	host, portStr, _ := net.SplitHostPort(listenAddr)
	port := 0
	fmt.Sscanf(portStr, "%d", &port)

	// Node parser map (mock)
	nodeParser := toki.ParserMap{
		toki.TypeNameOf(&iop.RegisterResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.RegisterResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeReadyResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeReadyResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	client, err := toki.DialTcp(ctx, host, port, 30, 10, nodeParser)
	if err != nil {
		t.Fatalf("failed to dial edge server: %v", err)
	}
	defer client.Close()

	// 3. Node가 RegisterRequest를 보내고 RegisterResponse를 받는다
	resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client.Communicator,
		&iop.RegisterRequest{Token: "test-token"},
		45*time.Second,
	)
	if err != nil {
		t.Fatalf("register request: %v", err)
	}
	if !resp.GetAccepted() {
		t.Fatalf("expected accepted register response, got reason %q", resp.GetReason())
	}
	rec, ok := nodeStore.FindByToken("test-token")
	if !ok {
		t.Fatal("expected record for test-token in store")
	}
	wantNodeID := rec.ID
	if wantNodeID != "node-test-01" {
		t.Fatalf("expected stored id to match explicit config id, got %q", wantNodeID)
	}
	if resp.GetNodeId() != wantNodeID {
		t.Fatalf("expected node id %q, got %q", wantNodeID, resp.GetNodeId())
	}
	if resp.GetAlias() != "test-node" {
		t.Fatalf("expected alias %q, got %q", "test-node", resp.GetAlias())
	}
	if resp.GetConfig().GetRuntime().GetConcurrency() != 2 {
		t.Fatalf("expected runtime concurrency 2, got %d", resp.GetConfig().GetRuntime().GetConcurrency())
	}
	if len(resp.GetConfig().GetAdapters()) != 2 {
		t.Fatalf("expected mock and ollama adapters, got %d", len(resp.GetConfig().GetAdapters()))
	}
	// Registration only claimed ownership and delivered config; drive the ready
	// signal so the connected lifecycle/event fire.
	signalReady(t, client, resp.GetNodeId())
	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeConnected {
			t.Fatalf("event type: got %q want %q", event.GetType(), eventpkg.TypeNodeConnected)
		}
		if event.GetNodeId() != wantNodeID || event.GetAlias() != "test-node" || event.GetReason() != eventpkg.ReasonRegistered {
			t.Fatalf("unexpected connected event: %+v", event)
		}
	case <-time.After(45 * time.Second):
		t.Fatal("timeout waiting for node connected event")
	}

	// 4. Registry 등록 여부 확인
	entry, ok := waitForRegistryEntry(ctx, registry, wantNodeID)
	if !ok {
		t.Fatal("node was not registered in edge registry within timeout")
	}

	if got := registry.Count(); got != 1 {
		t.Fatalf("expected registry count 1, got %d", got)
	}
	if entry == nil {
		t.Fatal("expected registry entry, got nil")
	}
	if got := entry.NodeID; got != wantNodeID {
		t.Fatalf("expected node id %q, got %q", wantNodeID, got)
	}
	if got := entry.Alias; got != "test-node" {
		t.Fatalf("expected alias %q, got %q", "test-node", got)
	}

	// 5. Alias로 Resolve 확인
	if e, err := registry.Resolve("test-node"); err != nil || e.NodeID != wantNodeID {
		t.Fatalf("failed to resolve by alias %q: %v", "test-node", err)
	}

	if err := client.Close(); err != nil {
		t.Fatalf("close client: %v", err)
	}
	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeDisconnected {
			t.Fatalf("event type: got %q want %q", event.GetType(), eventpkg.TypeNodeDisconnected)
		}
		if event.GetNodeId() != wantNodeID || event.GetAlias() != "test-node" || event.GetReason() != eventpkg.ReasonTransportClosed {
			t.Fatalf("unexpected disconnected event: %+v", event)
		}
		if event.GetMetadata()[eventpkg.MetadataTransportCloseReason] != toki.DisconnectReasonRemoteClosed {
			t.Fatalf("transport close reason: got %q want %q", event.GetMetadata()[eventpkg.MetadataTransportCloseReason], toki.DisconnectReasonRemoteClosed)
		}
		if event.GetMetadata()[eventpkg.MetadataTransportCloseError] == "" {
			t.Fatalf("expected transport close error metadata, got %+v", event.GetMetadata())
		}
	case <-time.After(45 * time.Second):
		t.Fatal("timeout waiting for node disconnected event")
	}
}

func TestEdgeServerRegistrationFailureReasons(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{
			ID:        "node-fail-01",
			Alias:     "fail-node",
			Token:     "valid-token",
			AgentKind: config.AgentKindGenericNode,
		},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}

	nodeEventCh := make(chan *iop.EdgeNodeEvent, 10)
	server.SetNodeEventHandler(func(event *iop.EdgeNodeEvent) {
		nodeEventCh <- event
	})

	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	// 1. Unknown Token Test
	client1 := dialNode(t, ctx, listenAddr)
	resp1, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client1.Communicator,
		&iop.RegisterRequest{Token: "unknown-token-12345"},
		45*time.Second,
	)
	if err != nil {
		t.Fatalf("send register 1: %v", err)
	}
	if resp1.GetAccepted() {
		t.Fatal("expected rejected")
	}

	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeRegistrationFailed {
			t.Fatalf("expected registration failed type, got %q", event.GetType())
		}
		if event.GetReason() != eventpkg.ReasonUnknownToken {
			t.Fatalf("expected reason unknown_token, got %q", event.GetReason())
		}
		if event.GetMetadata()[eventpkg.MetadataFailureReason] != eventpkg.ReasonUnknownToken {
			t.Fatalf("expected metadata failure reason, got %q", event.GetMetadata()[eventpkg.MetadataFailureReason])
		}
		if event.GetMetadata()[eventpkg.MetadataTokenPrefix] != "unknown-..." {
			t.Fatalf("expected token prefix unknown-..., got %q", event.GetMetadata()[eventpkg.MetadataTokenPrefix])
		}
	case <-time.After(45 * time.Second):
		t.Fatal("timeout waiting for unknown token registration failed event")
	}
	client1.Close()

	// 2. Duplicate Connection Test
	client2 := dialNode(t, ctx, listenAddr)
	resp2, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client2.Communicator,
		&iop.RegisterRequest{Token: "valid-token"},
		45*time.Second,
	)
	if err != nil {
		t.Fatalf("send register 2: %v", err)
	}
	if !resp2.GetAccepted() {
		t.Fatalf("expected register 2 accepted, got %q", resp2.GetReason())
	}
	// Ready signal opens the connected lifecycle/event for the accepted owner.
	signalReady(t, client2, resp2.GetNodeId())
	// Consume connection event
	<-nodeEventCh

	client3 := dialNode(t, ctx, listenAddr)
	resp3, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client3.Communicator,
		&iop.RegisterRequest{Token: "valid-token"},
		45*time.Second,
	)
	if err != nil {
		t.Fatalf("send register 3: %v", err)
	}
	if resp3.GetAccepted() {
		t.Fatal("expected duplicate connection registration rejected")
	}

	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeRegistrationFailed {
			t.Fatalf("expected registration failed type, got %q", event.GetType())
		}
		if event.GetReason() != eventpkg.ReasonDuplicateConnection {
			t.Fatalf("expected reason duplicate_connection, got %q", event.GetReason())
		}
		if event.GetMetadata()[eventpkg.MetadataFailureReason] != eventpkg.ReasonDuplicateConnection {
			t.Fatalf("expected failure reason metadata, got %q", event.GetMetadata()[eventpkg.MetadataFailureReason])
		}
		if event.GetMetadata()[eventpkg.MetadataAgentKind] != config.AgentKindGenericNode {
			t.Fatalf("expected agent kind metadata, got %q", event.GetMetadata()[eventpkg.MetadataAgentKind])
		}
	case <-time.After(45 * time.Second):
		t.Fatal("timeout waiting for duplicate registration failed event")
	}
	client2.Close()
	client3.Close()
}

func TestEdgeServerHeartbeatTimeoutDisconnect(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{
			ID:        "node-timeout-01",
			Alias:     "timeout-node",
			Token:     "timeout-token",
			AgentKind: config.AgentKindGenericNode,
		},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	server.HeartbeatInterval = 1
	server.HeartbeatWait = 1

	nodeEventCh := make(chan *iop.EdgeNodeEvent, 10)
	server.SetNodeEventHandler(func(event *iop.EdgeNodeEvent) {
		nodeEventCh <- event
	})

	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	// Register and ready first, while the client still answers heartbeats (the ready
	// signal opens the connected event), then stop answering so the server times the
	// connection out. Muting heartbeats before the handshake let HeartbeatWait=1 drop
	// the connection mid-register under the repeated -race suite ("not connected").
	client, resp := registerNode(t, ctx, listenAddr, "timeout-token")
	defer client.Close()
	if !resp.GetAccepted() {
		t.Fatalf("expected accepted, got %q", resp.GetReason())
	}
	client.RemoveListeners(toki.TypeNameOf(&packets.HeartBeat{}))
	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeConnected {
			t.Fatalf("expected connected event, got %q", event.GetType())
		}
	case <-time.After(45 * time.Second):
		t.Fatal("timeout waiting for connected event")
	}

	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeDisconnected {
			t.Fatalf("expected disconnected event, got %q", event.GetType())
		}
		if event.GetReason() != eventpkg.ReasonHeartbeatTimeout {
			t.Fatalf("expected reason heartbeat_timeout, got %q", event.GetReason())
		}
		if event.GetMetadata()[eventpkg.MetadataFailureReason] != eventpkg.ReasonHeartbeatTimeout {
			t.Fatalf("expected failure reason heartbeat_timeout, got %q", event.GetMetadata()[eventpkg.MetadataFailureReason])
		}
		if event.GetMetadata()[eventpkg.MetadataLifecycleState] != edgenode.LifecycleFailed {
			t.Fatalf("expected lifecycle state failed, got %q", event.GetMetadata()[eventpkg.MetadataLifecycleState])
		}
	case <-time.After(45 * time.Second):
		t.Fatal("timeout waiting for heartbeat timeout disconnected event")
	}
}

// disconnectRecorder captures the interleaving of the authoritative lifecycle
// callback and the observability event, which is the ordering under test, along
// with the connection generation the authoritative callback carries.
type disconnectRecorder struct {
	mu          sync.Mutex
	sequence    []string
	nodeIDs     []string
	reasons     []string
	generations []uint64
}

func (r *disconnectRecorder) onDisconnect(nodeID string, generation uint64, reason string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.sequence = append(r.sequence, "lifecycle")
	r.nodeIDs = append(r.nodeIDs, nodeID)
	r.reasons = append(r.reasons, reason)
	r.generations = append(r.generations, generation)
}

func (r *disconnectRecorder) onNodeEvent(event *iop.EdgeNodeEvent) {
	if event.GetType() != eventpkg.TypeNodeDisconnected {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.sequence = append(r.sequence, "event")
}

func (r *disconnectRecorder) snapshot() ([]string, []string, []string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]string(nil), r.sequence...),
		append([]string(nil), r.nodeIDs...),
		append([]string(nil), r.reasons...)
}

func (r *disconnectRecorder) generationSnapshot() []uint64 {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]uint64(nil), r.generations...)
}

func newDisconnectTestServer(t *testing.T, ctx context.Context, nodeID, token string) (*transport.Server, *edgenode.Registry, *disconnectRecorder, string) {
	t.Helper()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: nodeID, Alias: nodeID, Token: token, AgentKind: config.AgentKindGenericNode},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, zap.NewNop())
	if err != nil {
		t.Fatalf("new server: %v", err)
	}

	rec := &disconnectRecorder{}
	server.SetNodeDisconnectHandler(rec.onDisconnect)
	server.SetNodeEventHandler(rec.onNodeEvent)

	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	t.Cleanup(func() { server.Stop() })

	return server, registry, rec, listenAddr
}

// registerNode dials the edge and performs the full accepted-registration
// handshake resiliently: it retries dial+register on a transient connection
// stall (repeated -race suite) and, when accepted, drives the dispatch-ready
// signal so the node is fully connected (config claimed + ready) on return. A
// rejected registration is returned as-is with no ready signal. The returned
// client is the one that carried the successful register round-trip and must be
// used for any subsequent ready/disconnect steps.
func registerNode(t *testing.T, ctx context.Context, listenAddr, token string) (*toki.TcpClient, *iop.RegisterResponse) {
	t.Helper()
	client, resp := resilientRegister(t, ctx, listenAddr, token)
	if resp.GetAccepted() {
		signalReady(t, client, resp.GetNodeId())
	}
	return client, resp
}

// TestCurrentOwnerDisconnectReleasesLeasesBeforeEvent pins that the node whose
// connection actually owned the registry entry settles lifecycle cleanup exactly
// once, and that it happens before the disconnected event is published. The
// reverse order would let a dropped fanout strand the node's leases.
func TestCurrentOwnerDisconnectReleasesLeasesBeforeEvent(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	_, registry, rec, listenAddr := newDisconnectTestServer(t, ctx, "owner-01", "owner-token")

	client, resp := registerNode(t, ctx, listenAddr, "owner-token")
	if !resp.GetAccepted() {
		t.Fatalf("expected registration accepted, got %q", resp.GetReason())
	}
	if _, ok := waitForRegistryEntry(ctx, registry, "owner-01"); !ok {
		t.Fatal("registration not visible in registry")
	}

	if err := client.Close(); err != nil {
		t.Fatalf("close client: %v", err)
	}
	if ok := waitForRegistryAbsent(ctx, registry, "owner-01"); !ok {
		t.Fatal("registry still contains node after owner disconnect")
	}

	deadline := time.Now().Add(45 * time.Second)
	for {
		sequence, _, _ := rec.snapshot()
		if len(sequence) >= 2 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("timeout waiting for lifecycle+event, got %v", sequence)
		}
		time.Sleep(5 * time.Millisecond)
	}

	sequence, nodeIDs, reasons := rec.snapshot()
	if len(sequence) != 2 {
		t.Fatalf("expected exactly one lifecycle callback and one event, got %v", sequence)
	}
	if sequence[0] != "lifecycle" || sequence[1] != "event" {
		t.Errorf("expected cleanup before the observability event, got %v", sequence)
	}
	if len(nodeIDs) != 1 || nodeIDs[0] != "owner-01" {
		t.Errorf("lifecycle callback node ids = %v, want [owner-01]", nodeIDs)
	}
	if len(reasons) != 1 || reasons[0] == "" {
		t.Errorf("lifecycle callback reasons = %v, want one non-empty reason", reasons)
	}
}

// TestStaleDisconnectDoesNotReleaseLiveLease pins the ownership gate: a rejected
// duplicate connection closing must neither release the live connection's leases
// nor announce that the still-connected node disconnected.
func TestStaleDisconnectDoesNotReleaseLiveLease(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	_, registry, rec, listenAddr := newDisconnectTestServer(t, ctx, "stale-01", "stale-token")

	owner, resp := registerNode(t, ctx, listenAddr, "stale-token")
	defer owner.Close()
	if !resp.GetAccepted() {
		t.Fatalf("expected first registration accepted, got %q", resp.GetReason())
	}
	if _, ok := waitForRegistryEntry(ctx, registry, "stale-01"); !ok {
		t.Fatal("registration not visible in registry")
	}

	// The duplicate is rejected but still has a disconnect listener bound, which
	// is precisely the connection that must not touch the live entry.
	duplicate, dupResp := registerNode(t, ctx, listenAddr, "stale-token")
	if dupResp.GetAccepted() {
		t.Fatal("expected the duplicate registration to be rejected")
	}
	if err := duplicate.Close(); err != nil {
		t.Fatalf("close duplicate: %v", err)
	}

	// Give the stale disconnect listener room to misbehave before asserting.
	time.Sleep(200 * time.Millisecond)

	sequence, _, _ := rec.snapshot()
	if len(sequence) != 0 {
		t.Errorf("stale disconnect produced %v, want no lifecycle callback and no disconnected event", sequence)
	}
	if generations := rec.generationSnapshot(); len(generations) != 0 {
		t.Errorf("stale disconnect delivered generations %v, want none", generations)
	}
	if _, ok := registry.Get("stale-01"); !ok {
		t.Error("the live node was unregistered by the stale connection closing")
	}
}

// TestDisconnectHandlerReceivesCurrentGeneration pins that the authoritative
// disconnect callback carries the exact connection generation the registry
// assigned to the owner, so the service can fence that owner's leases by identity.
func TestDisconnectHandlerReceivesCurrentGeneration(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	_, registry, rec, listenAddr := newDisconnectTestServer(t, ctx, "gen-01", "gen-token")

	client, resp := registerNode(t, ctx, listenAddr, "gen-token")
	if !resp.GetAccepted() {
		t.Fatalf("expected registration accepted, got %q", resp.GetReason())
	}
	entry, ok := waitForRegistryEntry(ctx, registry, "gen-01")
	if !ok {
		t.Fatal("registration not visible in registry")
	}
	wantGeneration := entry.ConnectionGeneration
	if wantGeneration == 0 {
		t.Fatal("registered entry must carry a non-zero connection generation")
	}

	if err := client.Close(); err != nil {
		t.Fatalf("close client: %v", err)
	}
	if ok := waitForRegistryAbsent(ctx, registry, "gen-01"); !ok {
		t.Fatal("registry still contains node after owner disconnect")
	}

	deadline := time.Now().Add(45 * time.Second)
	for {
		if len(rec.generationSnapshot()) >= 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("timeout waiting for the authoritative disconnect callback")
		}
		time.Sleep(5 * time.Millisecond)
	}

	generations := rec.generationSnapshot()
	if len(generations) != 1 {
		t.Fatalf("expected exactly one lifecycle callback, got generations %v", generations)
	}
	if generations[0] != wantGeneration {
		t.Fatalf("disconnect generation = %d, want the owner's %d", generations[0], wantGeneration)
	}
}

// TestReconnectDisconnectGenerationsAreMonotonic pins that a node reconnecting on
// the same id draws a strictly higher generation and that each of its own
// disconnects reports that connection's generation, so an earlier connection's
// late callback can never be confused with the reconnect's identity.
func TestReconnectDisconnectGenerationsAreMonotonic(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	_, registry, rec, listenAddr := newDisconnectTestServer(t, ctx, "recon-01", "recon-token")

	connectAndCapture := func() uint64 {
		t.Helper()
		client, resp := registerNode(t, ctx, listenAddr, "recon-token")
		if !resp.GetAccepted() {
			t.Fatalf("expected registration accepted, got %q", resp.GetReason())
		}
		entry, ok := waitForRegistryEntry(ctx, registry, "recon-01")
		if !ok {
			t.Fatal("registration not visible in registry")
		}
		gen := entry.ConnectionGeneration
		if err := client.Close(); err != nil {
			t.Fatalf("close client: %v", err)
		}
		if ok := waitForRegistryAbsent(ctx, registry, "recon-01"); !ok {
			t.Fatal("registry still contains node after disconnect")
		}
		return gen
	}

	firstGen := connectAndCapture()
	secondGen := connectAndCapture()
	if secondGen <= firstGen {
		t.Fatalf("reconnect generation %d must exceed the first %d", secondGen, firstGen)
	}

	deadline := time.Now().Add(45 * time.Second)
	for {
		if len(rec.generationSnapshot()) >= 2 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("timeout waiting for two disconnect callbacks, got %v", rec.generationSnapshot())
		}
		time.Sleep(5 * time.Millisecond)
	}

	generations := rec.generationSnapshot()
	if len(generations) != 2 {
		t.Fatalf("expected two lifecycle callbacks, got %v", generations)
	}
	if generations[0] != firstGen || generations[1] != secondGen {
		t.Fatalf("disconnect generations = %v, want [%d %d]", generations, firstGen, secondGen)
	}
}

// TestHeartbeatTimeoutDisconnectCarriesCurrentGenerationBeforeEvent pins that a
// heartbeat-timeout disconnect settles the same authoritative boundary a normal
// close does: the lifecycle callback carries the registered owner's connection
// generation and reason=heartbeat_timeout, and it runs before the disconnected
// observability event. The reverse order (or a missing/zero generation) would let
// a dropped fanout strand the timed-out connection's leases or fence the wrong
// generation.
func TestHeartbeatTimeoutDisconnectCarriesCurrentGenerationBeforeEvent(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: "hb-01", Alias: "hb", Token: "hb-token", AgentKind: config.AgentKindGenericNode},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, zap.NewNop())
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	server.HeartbeatInterval = 1
	server.HeartbeatWait = 1

	rec := &disconnectRecorder{}
	server.SetNodeDisconnectHandler(rec.onDisconnect)
	server.SetNodeEventHandler(rec.onNodeEvent)

	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	// Register (and ready) first, while the client still answers heartbeats, so the
	// aggressive HeartbeatWait=1 cannot race the handshake and drop the connection
	// before it registers. Only after the entry is live do we stop answering
	// heartbeats, which is what drives the server to time this connection out. The
	// previous order — muting heartbeats before registering — let the server kill
	// the connection mid-handshake under the repeated -race suite ("not connected").
	client, resp := registerNode(t, ctx, listenAddr, "hb-token")
	defer client.Close()
	if !resp.GetAccepted() {
		t.Fatalf("expected registration accepted, got %q", resp.GetReason())
	}
	entry, ok := waitForRegistryEntry(ctx, registry, "hb-01")
	if !ok {
		t.Fatal("registration not visible in registry")
	}
	wantGeneration := entry.ConnectionGeneration
	if wantGeneration == 0 {
		t.Fatal("registered entry must carry a non-zero connection generation")
	}

	// Stop answering heartbeats so the server drives this connection to timeout.
	client.RemoveListeners(toki.TypeNameOf(&packets.HeartBeat{}))

	deadline := time.Now().Add(45 * time.Second)
	for {
		sequence, _, _ := rec.snapshot()
		if len(sequence) >= 2 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("timeout waiting for heartbeat lifecycle+event, got %v", sequence)
		}
		time.Sleep(5 * time.Millisecond)
	}

	sequence, nodeIDs, reasons := rec.snapshot()
	if len(sequence) != 2 {
		t.Fatalf("expected exactly one lifecycle callback and one event, got %v", sequence)
	}
	if sequence[0] != "lifecycle" || sequence[1] != "event" {
		t.Errorf("expected cleanup before the observability event, got %v", sequence)
	}
	if len(nodeIDs) != 1 || nodeIDs[0] != "hb-01" {
		t.Errorf("lifecycle callback node ids = %v, want [hb-01]", nodeIDs)
	}
	if len(reasons) != 1 || reasons[0] != eventpkg.ReasonHeartbeatTimeout {
		t.Errorf("lifecycle callback reasons = %v, want [%s]", reasons, eventpkg.ReasonHeartbeatTimeout)
	}
	generations := rec.generationSnapshot()
	if len(generations) != 1 {
		t.Fatalf("expected exactly one lifecycle callback, got generations %v", generations)
	}
	if generations[0] != wantGeneration {
		t.Fatalf("heartbeat disconnect generation = %d, want the owner's %d", generations[0], wantGeneration)
	}
}

// connectRecorder captures the interleaving of the authoritative connect
// callback and the observability connected event — the ordering under test —
// together with the generation each accepted connection carries. Only
// TypeNodeConnected events are recorded, so a rejected duplicate's
// registration-failed event and a close's disconnected event never appear.
type connectRecorder struct {
	mu          sync.Mutex
	sequence    []string
	nodeIDs     []string
	generations []uint64
}

func (r *connectRecorder) onConnect(nodeID string, generation uint64) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.sequence = append(r.sequence, "lifecycle")
	r.nodeIDs = append(r.nodeIDs, nodeID)
	r.generations = append(r.generations, generation)
}

func (r *connectRecorder) onNodeEvent(event *iop.EdgeNodeEvent) {
	if event.GetType() != eventpkg.TypeNodeConnected {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.sequence = append(r.sequence, "event")
}

func (r *connectRecorder) snapshot() ([]string, []uint64, []string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]string(nil), r.sequence...),
		append([]uint64(nil), r.generations...),
		append([]string(nil), r.nodeIDs...)
}

// TestAcceptedConnectHookRunsBeforeEventPerGeneration pins the accepted-connect
// authoritative lifecycle hook: an accepted registration fires exactly one
// connect callback carrying that connection's generation, before the connected
// event is published, and a reconnect on the same id fires a second callback with
// a strictly higher generation. A rejected duplicate fires no callback at all, so
// only genuine current owners drive resource activation.
func TestAcceptedConnectHookRunsBeforeEventPerGeneration(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: "connect-01", Alias: "connect", Token: "connect-token", AgentKind: config.AgentKindGenericNode},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, zap.NewNop())
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	rec := &connectRecorder{}
	server.SetNodeConnectHandler(rec.onConnect)
	server.SetNodeEventHandler(rec.onNodeEvent)
	if !server.HasNodeConnectHandler() {
		t.Fatal("connect handler was not registered")
	}
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	// First accepted registration.
	first, firstResp := registerNode(t, ctx, listenAddr, "connect-token")
	if !firstResp.GetAccepted() {
		t.Fatalf("expected first registration accepted, got %q", firstResp.GetReason())
	}
	firstEntry, ok := waitForRegistryEntry(ctx, registry, "connect-01")
	if !ok {
		t.Fatal("first registration not visible in registry")
	}
	firstGen := firstEntry.ConnectionGeneration
	if firstGen == 0 {
		t.Fatal("accepted entry must carry a non-zero connection generation")
	}

	// A duplicate registration while the owner is still connected is rejected and
	// must not fire the connect callback: only accepted current owners activate.
	duplicate, dupResp := registerNode(t, ctx, listenAddr, "connect-token")
	if dupResp.GetAccepted() {
		t.Fatal("expected the duplicate registration to be rejected")
	}
	if err := duplicate.Close(); err != nil {
		t.Fatalf("close duplicate: %v", err)
	}

	// Disconnect the owner, then reconnect the same id.
	if err := first.Close(); err != nil {
		t.Fatalf("close first client: %v", err)
	}
	if ok := waitForRegistryAbsent(ctx, registry, "connect-01"); !ok {
		t.Fatal("registry still contains node after disconnect")
	}
	second, secondResp := registerNode(t, ctx, listenAddr, "connect-token")
	defer second.Close()
	if !secondResp.GetAccepted() {
		t.Fatalf("expected reconnect accepted, got %q", secondResp.GetReason())
	}
	secondEntry, ok := waitForRegistryEntry(ctx, registry, "connect-01")
	if !ok {
		t.Fatal("reconnect registration not visible in registry")
	}
	secondGen := secondEntry.ConnectionGeneration
	if secondGen <= firstGen {
		t.Fatalf("reconnect generation %d must exceed the first %d", secondGen, firstGen)
	}

	deadline := time.Now().Add(45 * time.Second)
	for {
		if _, gens, _ := rec.snapshot(); len(gens) >= 2 {
			break
		}
		if time.Now().After(deadline) {
			_, gens, _ := rec.snapshot()
			t.Fatalf("timeout waiting for two connect callbacks, got %v", gens)
		}
		time.Sleep(5 * time.Millisecond)
	}

	sequence, generations, nodeIDs := rec.snapshot()
	// Exactly two accepted connections (first + reconnect); the rejected duplicate
	// contributed no callback.
	if len(generations) != 2 {
		t.Fatalf("expected exactly two connect callbacks (accept + reconnect), got generations %v", generations)
	}
	if generations[0] != firstGen || generations[1] != secondGen {
		t.Fatalf("connect generations = %v, want [%d %d]", generations, firstGen, secondGen)
	}
	for _, id := range nodeIDs {
		if id != "connect-01" {
			t.Fatalf("connect callback node ids = %v, want all connect-01", nodeIDs)
		}
	}
	// Each accepted connection's lifecycle callback precedes its connected event,
	// and no rejected/disconnected event interleaves.
	want := []string{"lifecycle", "event", "lifecycle", "event"}
	if len(sequence) != len(want) {
		t.Fatalf("connect/event sequence = %v, want %v", sequence, want)
	}
	for i := range want {
		if sequence[i] != want[i] {
			t.Fatalf("connect/event sequence = %v, want %v", sequence, want)
		}
	}
}

// TestRegistrationDoesNotDispatchBeforeReady pins the core fix: an accepted
// registration claims ownership and delivers config but fires no connect
// lifecycle and no connected event, and the entry is absent from dispatch-ready
// lookups. Only the node's NodeReadyRequest opens dispatch and, at that point,
// runs the connect lifecycle and connected event exactly once.
func TestRegistrationDoesNotDispatchBeforeReady(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: "gate-01", Alias: "gate", Token: "gate-token", AgentKind: config.AgentKindGenericNode},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}
	server, err := transport.NewServer(listenAddr, registry, nodeStore, zap.NewNop())
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	rec := &connectRecorder{}
	server.SetNodeConnectHandler(rec.onConnect)
	server.SetNodeEventHandler(rec.onNodeEvent)
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	// Register only (no ready): this test asserts the pending state before it later
	// signals ready. resilientRegister retries the dial+register on a transient
	// connection stall under the repeated -race suite.
	client, resp := resilientRegister(t, ctx, listenAddr, "gate-token")
	defer client.Close()
	if !resp.GetAccepted() {
		t.Fatalf("expected accepted, got %q", resp.GetReason())
	}
	if resp.GetConfig() == nil {
		t.Fatal("registration must still deliver config")
	}

	entry, ok := waitForRegistryEntry(ctx, registry, "gate-01")
	if !ok {
		t.Fatal("accepted registration must claim the registry id")
	}
	if entry.DispatchReady {
		t.Fatal("entry must be pending, not dispatch-ready, before the ready signal")
	}

	// Before ready: no connect lifecycle, no connected event, not dispatchable.
	time.Sleep(200 * time.Millisecond)
	if seq, _, _ := rec.snapshot(); len(seq) != 0 {
		t.Fatalf("registration fired connect/event before ready: %v", seq)
	}
	if n := len(registry.AllReady()); n != 0 {
		t.Fatalf("pending entry must not be dispatch-ready, AllReady=%d", n)
	}

	// Ready opens dispatch: exactly one lifecycle then one connected event.
	signalReady(t, client, resp.GetNodeId())
	deadline := time.Now().Add(45 * time.Second)
	for {
		if seq, _, _ := rec.snapshot(); len(seq) >= 2 {
			break
		}
		if time.Now().After(deadline) {
			seq, _, _ := rec.snapshot()
			t.Fatalf("timeout waiting for ready lifecycle+event, got %v", seq)
		}
		time.Sleep(5 * time.Millisecond)
	}
	seq, gens, nodeIDs := rec.snapshot()
	if len(seq) != 2 || seq[0] != "lifecycle" || seq[1] != "event" {
		t.Fatalf("ready sequence = %v, want [lifecycle event]", seq)
	}
	if len(gens) != 1 || gens[0] != entry.ConnectionGeneration {
		t.Fatalf("ready generation = %v, want [%d]", gens, entry.ConnectionGeneration)
	}
	if len(nodeIDs) != 1 || nodeIDs[0] != "gate-01" {
		t.Fatalf("ready node ids = %v, want [gate-01]", nodeIDs)
	}
	if n := len(registry.AllReady()); n != 1 {
		t.Fatalf("ready entry must be dispatchable, AllReady=%d", n)
	}
}

// TestReadyHandshakeIdempotentAndStaleRejected pins that a duplicate ready for
// the current owner acks success without repeating the connect lifecycle/event,
// and that a ready from a non-owner (a rejected duplicate connection) is rejected
// and drives no lifecycle at all.
func TestReadyHandshakeIdempotentAndStaleRejected(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: "idem-01", Alias: "idem", Token: "idem-token", AgentKind: config.AgentKindGenericNode},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}
	server, err := transport.NewServer(listenAddr, registry, nodeStore, zap.NewNop())
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	rec := &connectRecorder{}
	server.SetNodeConnectHandler(rec.onConnect)
	server.SetNodeEventHandler(rec.onNodeEvent)
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	owner, resp := resilientRegister(t, ctx, listenAddr, "idem-token")
	defer owner.Close()
	if !resp.GetAccepted() {
		t.Fatalf("expected accepted, got %q", resp.GetReason())
	}

	// First ready: one lifecycle + one event.
	signalReady(t, owner, resp.GetNodeId())
	deadline := time.Now().Add(45 * time.Second)
	for {
		if seq, _, _ := rec.snapshot(); len(seq) >= 2 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("timeout waiting for first ready lifecycle+event")
		}
		time.Sleep(5 * time.Millisecond)
	}

	// Duplicate ready on the same owner: acked ready, but no extra lifecycle/event.
	dup, err := toki.SendRequestTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
		&owner.Communicator, &iop.NodeReadyRequest{NodeId: resp.GetNodeId()}, 45*time.Second)
	if err != nil {
		t.Fatalf("duplicate ready: %v", err)
	}
	if !dup.GetReady() {
		t.Fatalf("duplicate ready for current owner must ack ready, got reason %q", dup.GetReason())
	}

	// A rejected duplicate connection is not the owner; its ready must be rejected.
	second, dupReg := resilientRegister(t, ctx, listenAddr, "idem-token")
	defer second.Close()
	if dupReg.GetAccepted() {
		t.Fatal("expected duplicate registration to be rejected")
	}
	staleReady, err := toki.SendRequestTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
		&second.Communicator, &iop.NodeReadyRequest{NodeId: "idem-01"}, 45*time.Second)
	if err != nil {
		t.Fatalf("stale ready: %v", err)
	}
	if staleReady.GetReady() {
		t.Fatal("ready from a non-owner connection must be rejected")
	}

	// Still exactly one lifecycle + one connected event across all readys.
	time.Sleep(150 * time.Millisecond)
	seq, gens, _ := rec.snapshot()
	if len(seq) != 2 || seq[0] != "lifecycle" || seq[1] != "event" {
		t.Fatalf("duplicate/stale ready changed the sequence: %v, want [lifecycle event]", seq)
	}
	if len(gens) != 1 {
		t.Fatalf("expected exactly one connect callback, got generations %v", gens)
	}
}

// TestReadyDisconnectRace wires the real Service.HandleNodeConnect /
// HandleNodeDisconnect — the actual queue handlers, not a dummy — as the
// authoritative connect/disconnect lifecycle, and forces the ready/disconnect
// race deterministically. It preserves the old owner's entry snapshot at the
// ready transition, completes the real server-side disconnect listener, accepts
// a pending reconnect, then resumes the old continuation against the real queue
// handler and WithCurrentOwner event gate. The stale continuation activates
// nothing and emits no connected event; only the reconnect's own ready produces
// exactly one connected event.
//
// The deterministic proof that a stale generation neither pumps nor re-activates
// a queued waiter lives in the white-box service test
// TestReconnectActivationLinearizesAgainstDisconnect; here the same real handler
// is exercised across the transport ready/disconnect boundary.
func TestReadyDisconnectRace(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()

	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: "race-01", Alias: "race", Token: "race-token", AgentKind: config.AgentKindGenericNode},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}
	server, err := transport.NewServer(listenAddr, registry, nodeStore, zap.NewNop())
	if err != nil {
		t.Fatalf("new server: %v", err)
	}

	// The real queue handlers back the transport lifecycle.
	svc := edgeservice.New(registry, edgeevents.NewBus())

	var connectMu sync.Mutex
	var connectGenerations []uint64
	handleConnect := func(nodeID string, gen uint64) {
		connectMu.Lock()
		connectGenerations = append(connectGenerations, gen)
		connectMu.Unlock()
		// Actual queue handler: a superseded generation is a no-op here.
		svc.HandleNodeConnect(nodeID, gen)
	}
	server.SetNodeConnectHandler(handleConnect)
	disconnectDone := make(chan uint64, 1)
	server.SetNodeDisconnectHandler(func(nodeID string, generation uint64, reason string) {
		svc.HandleNodeDisconnect(nodeID, generation, reason)
		disconnectDone <- generation
	})

	var eventMutex sync.Mutex
	var connectedEvents []*iop.EdgeNodeEvent
	server.SetNodeEventHandler(func(e *iop.EdgeNodeEvent) {
		if e.GetType() != eventpkg.TypeNodeConnected {
			return
		}
		eventMutex.Lock()
		connectedEvents = append(connectedEvents, e)
		eventMutex.Unlock()
	})

	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	// 1. Old owner registers and becomes the current generation.
	client1, resp1 := resilientRegister(t, ctx, listenAddr, "race-token")
	defer client1.Close()
	if !resp1.GetAccepted() {
		t.Fatalf("register 1 rejected: %s", resp1.GetReason())
	}
	entry1, ok := registry.Get("race-01")
	if !ok {
		t.Fatal("expected race-01 in registry after registration")
	}
	oldGen := entry1.ConnectionGeneration
	client1ServerSide := entry1.Client

	// Preserve the exact snapshot that handleReadyRequest receives from the
	// pending→ready transition, but delay its connect/event continuation until
	// after the authoritative disconnect has settled.
	staleReadyEntry, transitioned, ok := registry.MarkDispatchReadyOwner("race-01", client1ServerSide)
	if !ok || !transitioned {
		t.Fatalf("old owner ready transition=(ok=%v transitioned=%v), want true,true", ok, transitioned)
	}

	// 2. Close the peer normally and wait for bindDisconnectListener to finish
	//    current-owner unregister plus the real Service queue settlement.
	if err := client1.Close(); err != nil {
		t.Fatalf("close old owner: %v", err)
	}
	select {
	case disconnectedGen := <-disconnectDone:
		if disconnectedGen != oldGen {
			t.Fatalf("disconnect generation=%d, want old owner %d", disconnectedGen, oldGen)
		}
	case <-time.After(45 * time.Second):
		t.Fatal("authoritative disconnect listener did not complete")
	}

	client2, resp2 := resilientRegister(t, ctx, listenAddr, "race-token")
	defer client2.Close()
	if !resp2.GetAccepted() {
		t.Fatalf("expected reconnect registration accepted, got reason: %s", resp2.GetReason())
	}
	entry2, ok := registry.Get("race-01")
	if !ok || entry2.Client == client1ServerSide {
		t.Fatal("registry must hold the reconnected owner")
	}
	if entry2.DispatchReady {
		t.Fatal("reconnected owner must remain pending until its own ready")
	}

	// 3. Resume the stale continuation. It runs the real queue handler but the
	// generation currency check makes it a no-op; the exact old entry snapshot is
	// also rejected by the same WithCurrentOwner gate used for connected events.
	handleConnect("race-01", oldGen)
	staleEventRan := registry.WithCurrentOwner(staleReadyEntry, func() {
		eventMutex.Lock()
		connectedEvents = append(connectedEvents, &iop.EdgeNodeEvent{Type: eventpkg.TypeNodeConnected})
		eventMutex.Unlock()
	})
	if staleEventRan {
		t.Fatal("stale ready snapshot passed the connected-event owner gate")
	}
	eventMutex.Lock()
	staleEvents := len(connectedEvents)
	eventMutex.Unlock()
	if staleEvents != 0 {
		t.Fatalf("stale generation emitted %d connected events, want 0", staleEvents)
	}

	entryAfter, _ := registry.Get("race-01")
	if entryAfter.DispatchReady {
		t.Error("reconnected owner must remain pending after the stale continuation")
	}

	// 4. The reconnect's own ready activates and emits exactly one connected event.
	signalReady(t, client2, resp2.GetNodeId())
	readyDeadline := time.Now().Add(45 * time.Second)
	for {
		eventMutex.Lock()
		n := len(connectedEvents)
		eventMutex.Unlock()
		if n >= 1 {
			break
		}
		if time.Now().After(readyDeadline) {
			t.Fatal("reconnect ready did not emit a connected event")
		}
		time.Sleep(5 * time.Millisecond)
	}
	eventMutex.Lock()
	numEvents := len(connectedEvents)
	eventMutex.Unlock()
	if numEvents != 1 {
		t.Errorf("expected exactly 1 connected event after reconnect ready, got %d", numEvents)
	}

	// The real connect handler was invoked for the reconnected generation too.
	connectMu.Lock()
	sawNew := false
	for _, g := range connectGenerations {
		if g == entry2.ConnectionGeneration {
			sawNew = true
		}
	}
	connectMu.Unlock()
	if !sawNew {
		t.Fatal("real connect handler was not invoked for the reconnected generation")
	}
}