iop/apps/node/internal/bootstrap/module_test.go

package bootstrap_test

import (
	"context"
	"errors"
	"fmt"
	"net"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"git.toki-labs.com/toki/proto-socket/go/packets"
	"go.uber.org/fx"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

	"iop/apps/node/internal/bootstrap"
	"iop/apps/node/internal/transport"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

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

// bindEdgeReadyResponder answers the node's post-handler NodeReadyRequest with a
// positive ack, so connectRuntime's SignalReady completes and the session is
// treated as fully established. Mock edges that accept registration must install
// it or the node's ready handshake times out and reconnects. When readyCh is
// non-nil it fires only after the ready response has been written to the socket.
// This lets a test safely close the edge side after a fully established session,
// instead of racing the in-flight ready handshake.
func bindEdgeReadyResponder(client *toki.TcpClient, readyCh chan<- struct{}) {
	client.Communicator.AddRequestListener(
		toki.TypeNameOf(&iop.NodeReadyRequest{}),
		func(message proto.Message, requestNonce int32) {
			if _, ok := message.(*iop.NodeReadyRequest); !ok {
				return
			}
			data, err := proto.Marshal(&iop.NodeReadyResponse{Ready: true})
			if err != nil {
				return
			}
			if err := client.QueuePacket(&packets.PacketBase{
				TypeName:      toki.TypeNameOf(&iop.NodeReadyResponse{}),
				ResponseNonce: requestNonce,
				Data:          data,
			}); err != nil {
				return
			}
			if readyCh != nil {
				select {
				case readyCh <- struct{}{}:
				default:
				}
			}
		},
	)
}

func freeAddrForBootstrap(t *testing.T) (host string, port int, addr string) {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	addr = l.Addr().String()
	l.Close()
	h, portStr, _ := net.SplitHostPort(addr)
	fmt.Sscanf(portStr, "%d", &port)
	return h, port, addr
}

func useTempCwd(t *testing.T) {
	t.Helper()
	previous, err := os.Getwd()
	if err != nil {
		t.Fatalf("get cwd: %v", err)
	}
	if err := os.Chdir(t.TempDir()); err != nil {
		t.Fatalf("chdir temp: %v", err)
	}
	t.Cleanup(func() {
		if err := os.Chdir(previous); err != nil {
			t.Fatalf("restore cwd: %v", err)
		}
	})
}

// startAcceptingEdge starts a mock edge bound to host:port that accepts every
// RegisterRequest. The returned channels fire once per registration and
// disconnect, respectively.
func startAcceptingEdge(t *testing.T, ctx context.Context, host string, port int, nodeID string) (chan struct{}, chan struct{}) {
	t.Helper()
	regCh := make(chan struct{}, 16)
	disconnectCh := make(chan struct{}, 16)
	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
		client.AddDisconnectListener(func(_ *toki.TcpClient) {
			disconnectCh <- struct{}{}
		})
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				regCh <- struct{}{}
				return &iop.RegisterResponse{
					Accepted: true,
					NodeId:   nodeID,
					Config:   &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
				}, nil
			},
		)
		bindEdgeReadyResponder(client, nil)
		return client
	})
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start accepting edge: %v", err)
	}
	t.Cleanup(func() { server.Stop() })
	return regCh, disconnectCh
}

// startRejectingEdge starts a mock edge bound to host:port that rejects every
// RegisterRequest with the given reason.
func startRejectingEdge(t *testing.T, ctx context.Context, host string, port int, reason string) {
	t.Helper()
	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				return &iop.RegisterResponse{Accepted: false, Reason: reason}, nil
			},
		)
		return client
	})
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start rejecting edge: %v", err)
	}
	t.Cleanup(func() { server.Stop() })
}

// TestSupervisorConnectsWhenEdgeStartsAfterNode verifies the actual DialEdge
// connection-refused path: Node starts first against a closed port, keeps
// retrying in the same process, and establishes exactly one session after Edge
// begins listening on that same port. SDD S16.
func TestSupervisorConnectsWhenEdgeStartsAfterNode(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	const refusedAttempts = 3
	var dialCount int32
	retryReady := make(chan struct{}, refusedAttempts)
	releaseRetry := make(chan struct{})
	observedDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		atomic.AddInt32(&dialCount, 1)
		return transport.DialEdge(ctx, edgeAddr, token, logger)
	})
	fakeSleeper := func(ctx context.Context, _ time.Duration) {
		retryReady <- struct{}{}
		if atomic.LoadInt32(&dialCount) >= refusedAttempts {
			select {
			case <-releaseRetry:
			case <-ctx.Done():
			}
		}
	}

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(observedDialer), bootstrap.WithSleeper(fakeSleeper)), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()
	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	// Edge is not listening yet. Each completed sleep entry follows a real
	// DialEdge connection-refused result. Hold the third retry until the listener
	// is live so the next actual dial succeeds deterministically.
	for attempt := 1; attempt <= refusedAttempts; attempt++ {
		select {
		case <-retryReady:
		case <-time.After(5 * time.Second):
			t.Fatalf("timeout waiting for refused DialEdge attempt %d", attempt)
		}
	}
	if count := atomic.LoadInt32(&dialCount); count != refusedAttempts {
		t.Fatalf("expected %d real refused dials before Edge startup, got %d", refusedAttempts, count)
	}

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)
	regCh, _ := startAcceptingEdge(t, serverCtx, host, port, "delayed-node")
	close(releaseRetry)

	select {
	case <-regCh:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for delayed initial registration")
	}

	if count := atomic.LoadInt32(&dialCount); count != refusedAttempts+1 {
		t.Fatalf("expected %d dials (3 real refusals + 1 success), got %d", refusedAttempts+1, count)
	}

	// Exactly one session must be established.
	select {
	case <-regCh:
		t.Fatal("unexpected second registration; supervisor should hold one session")
	case <-time.After(200 * time.Millisecond):
	}
}

// TestSupervisorUnlimitedAttemptsPassesTen verifies that an explicit
// max_attempts=0 (unlimited) keeps retrying past the former hard-coded 10-attempt
// cap and eventually connects. SDD S16/S17.
func TestSupervisorUnlimitedAttemptsPassesTen(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)
	regCh, _ := startAcceptingEdge(t, serverCtx, host, port, "unlimited-node")

	const failFirst = 15 // exceeds the old hard-coded 10-attempt cap
	var dialCount int32
	fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		if atomic.AddInt32(&dialCount, 1) <= failFirst {
			return nil, errors.New("edge unavailable")
		}
		return transport.DialEdge(ctx, edgeAddr, token, logger)
	})
	fakeSleeper := func(context.Context, time.Duration) {}

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper)), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()
	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	select {
	case <-regCh:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for registration after >10 unlimited attempts")
	}

	count := atomic.LoadInt32(&dialCount)
	if count != failFirst+1 {
		t.Fatalf("expected %d dials, got %d", failFirst+1, count)
	}
	if count <= 10 {
		t.Fatalf("unlimited policy must retry past 10 attempts, got %d", count)
	}
}

// TestSupervisorFiniteExhaustionExitsOne verifies that a finite reconnect policy
// makes exactly max_attempts initial-connect attempts and then terminates the
// node with exit code 1. SDD S17.
func TestSupervisorFiniteExhaustionExitsOne(t *testing.T) {
	useTempCwd(t)
	_, _, addr := freeAddrForBootstrap(t) // no server listening

	var dialCount int32
	fakeDialer := bootstrap.DialFunc(func(_ context.Context, _, _ string, _ *zap.Logger) (*transport.RegisterResult, error) {
		atomic.AddInt32(&dialCount, 1)
		return nil, errors.New("edge unavailable")
	})
	fakeSleeper := func(context.Context, time.Duration) {}

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 3, IntervalSec: 0},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper)), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()
	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	select {
	case sig := <-app.Wait():
		if sig.ExitCode != 1 {
			t.Fatalf("expected exit code 1, got %d", sig.ExitCode)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for finite exhaustion shutdown")
	}

	if count := atomic.LoadInt32(&dialCount); count != 3 {
		t.Fatalf("expected exactly 3 initial-connect attempts, got %d", count)
	}
}

// TestSupervisorFatalFailureDoesNotRetry verifies that an authoritative
// registration rejection is non-retryable: the supervisor exits with code 1
// after a single attempt without any retry. SDD S17.
func TestSupervisorFatalFailureDoesNotRetry(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)
	startRejectingEdge(t, serverCtx, host, port, "bad-token")

	var dialCount int32
	fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		atomic.AddInt32(&dialCount, 1)
		return transport.DialEdge(ctx, edgeAddr, token, logger)
	})

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 5, IntervalSec: 0},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer)), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()
	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	select {
	case sig := <-app.Wait():
		if sig.ExitCode != 1 {
			t.Fatalf("expected exit code 1 on fatal failure, got %d", sig.ExitCode)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for fatal shutdown")
	}

	if count := atomic.LoadInt32(&dialCount); count != 1 {
		t.Fatalf("fatal failure must not retry: expected 1 dial, got %d", count)
	}
}

// TestSupervisorShutdownCancelsRetry verifies that a local shutdown cancels a
// pending retry sleep, ends cleanly with no error, and runs no additional
// attempt. SDD S17.
func TestSupervisorShutdownCancelsRetry(t *testing.T) {
	useTempCwd(t)
	_, _, addr := freeAddrForBootstrap(t) // no server listening

	var dialCount int32
	fakeDialer := bootstrap.DialFunc(func(_ context.Context, _, _ string, _ *zap.Logger) (*transport.RegisterResult, error) {
		atomic.AddInt32(&dialCount, 1)
		return nil, errors.New("edge unavailable")
	})

	sleeping := make(chan struct{}, 1)
	fakeSleeper := func(ctx context.Context, _ time.Duration) {
		select {
		case sleeping <- struct{}{}:
		default:
		}
		<-ctx.Done() // block until the supervisor is cancelled by shutdown
	}

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper)), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()
	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}

	// Wait until the supervisor has dialed once and parked in the retry sleep.
	select {
	case <-sleeping:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for supervisor to enter retry sleep")
	}
	countAtSleep := atomic.LoadInt32(&dialCount)
	if countAtSleep != 1 {
		t.Fatalf("expected exactly 1 dial before parking in sleep, got %d", countAtSleep)
	}

	// Local shutdown must cancel the sleep and stop cleanly with no error.
	stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer stopCancel()
	if err := app.Stop(stopCtx); err != nil {
		t.Fatalf("app stop must be clean, got: %v", err)
	}

	if count := atomic.LoadInt32(&dialCount); count != countAtSleep {
		t.Fatalf("shutdown must not trigger extra dial: had %d, got %d", countAtSleep, count)
	}
}

// TestSupervisorDoesNotStartMetricsWhenPortDisabled guards the Node default
// against treating port zero as an ephemeral wildcard listener. It waits for
// the reconnect supervisor to be active before asserting the injected starter
// was never called, then verifies normal shutdown still cancels the retry.
func TestSupervisorDoesNotStartMetricsWhenPortDisabled(t *testing.T) {
	useTempCwd(t)
	_, _, addr := freeAddrForBootstrap(t)

	fakeDialer := bootstrap.DialFunc(func(_ context.Context, _, _ string, _ *zap.Logger) (*transport.RegisterResult, error) {
		return nil, errors.New("edge unavailable")
	})
	sleeping := make(chan struct{}, 1)
	fakeSleeper := func(ctx context.Context, _ time.Duration) {
		select {
		case sleeping <- struct{}{}:
		default:
		}
		<-ctx.Done()
	}
	metricsStarted := make(chan struct{}, 1)
	metricsStarter := func(int) error {
		metricsStarted <- struct{}{}
		return nil
	}

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
		Logging:   config.LoggingConf{Level: "error"},
		Metrics:   config.MetricsConf{Port: 0},
	}
	app := fx.New(
		bootstrap.Module(
			cfg,
			bootstrap.WithDialer(fakeDialer),
			bootstrap.WithSleeper(fakeSleeper),
			bootstrap.WithMetricsStarter(metricsStarter),
		),
		fx.NopLogger,
	)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()
	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}

	select {
	case <-sleeping:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for supervisor to enter retry sleep")
	}
	select {
	case <-metricsStarted:
		t.Fatal("metrics starter must not run when metrics port is disabled")
	case <-time.After(100 * time.Millisecond):
	}

	stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer stopCancel()
	if err := app.Stop(stopCtx); err != nil {
		t.Fatalf("app stop must be clean, got: %v", err)
	}
}

// TestSupervisorShutdownWinsInFlightDialFatal fixes the ordering between a
// local stop and a dialer that observes cancellation but returns a fatal result
// afterward. Cancellation must produce a clean stop, no exit 1, and no retry.
func TestSupervisorShutdownWinsInFlightDialFatal(t *testing.T) {
	useTempCwd(t)
	_, _, addr := freeAddrForBootstrap(t)

	started := make(chan struct{}, 1)
	var dialCount int32
	dialer := bootstrap.DialFunc(func(ctx context.Context, _, _ string, logger *zap.Logger) (*transport.RegisterResult, error) {
		atomic.AddInt32(&dialCount, 1)
		started <- struct{}{}
		<-ctx.Done()
		return transport.DialEdge(context.Background(), "127.0.0.1:0", "tok", logger)
	})
	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(dialer)), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()
	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	select {
	case <-started:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for in-flight dial")
	}

	stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer stopCancel()
	if err := app.Stop(stopCtx); err != nil {
		t.Fatalf("app stop must be clean, got: %v", err)
	}
	if count := atomic.LoadInt32(&dialCount); count != 1 {
		t.Fatalf("shutdown must not retry fatal late result: got %d dials", count)
	}
	select {
	case sig := <-app.Wait():
		t.Fatalf("shutdown must not request failure exit, got %+v", sig)
	case <-time.After(100 * time.Millisecond):
	}
}

// TestSupervisorShutdownClosesInFlightDialSuccess verifies that a dialer which
// observes cancellation and then returns a successful registered session cannot
// install or leak that late owner. The one registered session is closed by the
// post-dial cancellation fence and no failure exit or retry is emitted.
func TestSupervisorShutdownClosesInFlightDialSuccess(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)
	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)
	regCh, disconnectCh := startAcceptingEdge(t, serverCtx, host, port, "late-success-node")

	started := make(chan struct{}, 1)
	var dialCount int32
	dialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		atomic.AddInt32(&dialCount, 1)
		started <- struct{}{}
		<-ctx.Done()
		lateCtx, lateCancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer lateCancel()
		return transport.DialEdge(lateCtx, edgeAddr, token, logger)
	})
	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(dialer)), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()
	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	select {
	case <-started:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for in-flight dial")
	}

	stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer stopCancel()
	stopDone := make(chan error, 1)
	go func() { stopDone <- app.Stop(stopCtx) }()

	select {
	case <-regCh:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for late successful registration")
	}
	select {
	case err := <-stopDone:
		if err != nil {
			t.Fatalf("app stop must be clean, got: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for app stop")
	}
	select {
	case <-disconnectCh:
	case <-time.After(5 * time.Second):
		t.Fatal("late successful session leaked after supervisor shutdown")
	}
	if count := atomic.LoadInt32(&dialCount); count != 1 {
		t.Fatalf("shutdown must not retry late success: got %d dials", count)
	}
	select {
	case <-regCh:
		t.Fatal("unexpected second registration after shutdown")
	case sig := <-app.Wait():
		t.Fatalf("shutdown must not request failure exit, got %+v", sig)
	case <-time.After(100 * time.Millisecond):
	}
}

// TestReconnectSupervisorReestablishesSession verifies that after the edge-side
// client closes the connection, the supervisor reconnects and re-registers.
func TestReconnectSupervisorReestablishesSession(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)

	// regCh fires each time a RegisterRequest is handled; readyCh fires when a
	// connection completes its ready handshake so the test can close it only after
	// the session is fully established.
	regCh := make(chan struct{}, 8)
	readyCh := make(chan struct{}, 8)
	acceptedCh := make(chan *toki.TcpClient, 8)

	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				regCh <- struct{}{}
				return &iop.RegisterResponse{
					Accepted: true,
					NodeId:   "reconnect-test-node",
					Config:   &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
				}, nil
			},
		)
		bindEdgeReadyResponder(client, readyCh)
		acceptedCh <- client
		return client
	})
	if err := server.Start(serverCtx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	t.Cleanup(func() { server.Stop() })

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 5, IntervalSec: 0},
		Logging:   config.LoggingConf{Level: "error"},
	}

	// Wrap the real dialer so the test can wait for the initial session to be
	// fully established before inducing a disconnect (OnStart is non-blocking).
	established := make(chan struct{}, 1)
	var establishOnce sync.Once
	fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		res, err := transport.DialEdge(ctx, edgeAddr, token, logger)
		if err == nil {
			establishOnce.Do(func() { established <- struct{}{} })
		}
		return res, err
	})

	app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer)), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()

	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	// Wait for first registration and full establishment.
	select {
	case <-regCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first registration")
	}
	select {
	case <-established:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for initial session establishment")
	}
	var firstClient *toki.TcpClient
	select {
	case firstClient = <-acceptedCh:
	case <-time.After(time.Second):
		t.Fatal("timeout waiting for first accepted client")
	}
	// Wait for the ready handshake before closing so the disconnect is induced on a
	// fully-established session, not mid-SignalReady.
	select {
	case <-readyCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for initial ready handshake")
	}

	// Force edge-side close → remote disconnect on node.
	_ = firstClient.Close()

	// Wait for second registration (supervisor reconnected and re-registered).
	select {
	case <-regCh:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for supervisor reconnect registration")
	}
}

// TestReconnectSupervisorExhaustsRetries verifies that when all reconnect
// attempts fail, the node cleans up resources and calls Shutdown(ExitCode(1)).
func TestReconnectSupervisorExhaustsRetries(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)

	acceptedCh := make(chan *toki.TcpClient, 4)
	readyCh := make(chan struct{}, 4)
	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				return &iop.RegisterResponse{
					Accepted: true,
					NodeId:   "exhaust-test-node",
					Config:   &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
				}, nil
			},
		)
		bindEdgeReadyResponder(client, readyCh)
		acceptedCh <- client
		return client
	})
	if err := server.Start(serverCtx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	t.Cleanup(func() { server.Stop() })

	// Fake dialer: first call uses real DialEdge; subsequent calls fail immediately.
	// established fires once the initial session is fully established so the test
	// can induce the disconnect without racing the initial registration.
	established := make(chan struct{}, 1)
	var dialCount int32
	fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		if atomic.AddInt32(&dialCount, 1) == 1 {
			res, err := transport.DialEdge(ctx, edgeAddr, token, logger)
			if err == nil {
				established <- struct{}{}
			}
			return res, err
		}
		return nil, errors.New("fake reconnect failure")
	})

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 3, IntervalSec: 0},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(
		bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer)),
		fx.NopLogger,
	)

	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()

	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	// Wait for the initial session to establish, then close it from the edge side.
	select {
	case <-established:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for initial session establishment")
	}
	var firstClient *toki.TcpClient
	select {
	case firstClient = <-acceptedCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first connection")
	}
	// Close only after the ready handshake so the first session counts as an
	// established connection, not a mid-SignalReady failure.
	select {
	case <-readyCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for initial ready handshake")
	}
	_ = firstClient.Close()

	// Supervisor exhausts retries → Shutdown(ExitCode(1)).
	select {
	case sig := <-app.Wait():
		if sig.ExitCode != 1 {
			t.Fatalf("expected exit code 1, got %d", sig.ExitCode)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for supervisor exhaustion shutdown")
	}

	// 1 initial dial + 3 retry attempts = 4 total.
	if count := atomic.LoadInt32(&dialCount); count != 4 {
		t.Fatalf("expected 4 dial attempts (1 initial + 3 retries), got %d", count)
	}
}

// TestReconnectSupervisorPolicyTimingAndLimit verifies that the supervisor
// respects MaxAttempts=10 and that the fake sleeper receives IntervalSec=10s
// before each reconnect attempt (SDD reconnect_waiting state).
func TestReconnectSupervisorPolicyTimingAndLimit(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)

	acceptedCh := make(chan *toki.TcpClient, 4)
	readyCh := make(chan struct{}, 4)
	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				return &iop.RegisterResponse{
					Accepted: true,
					NodeId:   "policy-test-node",
					Config:   &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
				}, nil
			},
		)
		bindEdgeReadyResponder(client, readyCh)
		acceptedCh <- client
		return client
	})
	if err := server.Start(serverCtx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	t.Cleanup(func() { server.Stop() })

	// Fake dialer: first call uses real DialEdge; subsequent calls fail immediately.
	// established fires once the initial session is fully established so the test
	// can induce the disconnect without racing the initial registration.
	established := make(chan struct{}, 1)
	var dialCount int32
	fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		if atomic.AddInt32(&dialCount, 1) == 1 {
			res, err := transport.DialEdge(ctx, edgeAddr, token, logger)
			if err == nil {
				established <- struct{}{}
			}
			return res, err
		}
		return nil, errors.New("fake reconnect failure")
	})

	// Fake sleeper records received durations without blocking.
	var slMu sync.Mutex
	var sleptDurations []time.Duration
	fakeSleeper := func(_ context.Context, d time.Duration) {
		slMu.Lock()
		sleptDurations = append(sleptDurations, d)
		slMu.Unlock()
	}

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 10, IntervalSec: 10},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(
		bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper)),
		fx.NopLogger,
	)

	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()

	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	// Wait for the initial session to establish, then close it from the edge side.
	select {
	case <-established:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for initial session establishment")
	}
	var firstClient *toki.TcpClient
	select {
	case firstClient = <-acceptedCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first connection")
	}
	// Close only after the ready handshake so the first session counts as an
	// established connection, not a mid-SignalReady failure.
	select {
	case <-readyCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for initial ready handshake")
	}
	_ = firstClient.Close()

	// Supervisor exhausts all 10 retries → Shutdown(ExitCode(1)).
	select {
	case sig := <-app.Wait():
		if sig.ExitCode != 1 {
			t.Fatalf("expected exit code 1, got %d", sig.ExitCode)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for supervisor exhaustion shutdown")
	}

	// 1 initial dial + 10 retry attempts = 11 total.
	if count := atomic.LoadInt32(&dialCount); count != 11 {
		t.Fatalf("expected 11 dial attempts (1 initial + 10 retries), got %d", count)
	}

	// Each retry is preceded by a sleep of IntervalSec=10s (reconnect_waiting state).
	slMu.Lock()
	durations := append([]time.Duration(nil), sleptDurations...)
	slMu.Unlock()
	if len(durations) != 10 {
		t.Fatalf("expected 10 sleep calls (one before each retry), got %d: %v", len(durations), durations)
	}
	for i, d := range durations {
		if d != 10*time.Second {
			t.Errorf("sleep[%d] = %v, want 10s", i, d)
		}
	}
}