iop/apps/node/internal/transport/session_test.go

package transport_test

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

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

	"iop/apps/node/internal/transport"
	iop "iop/proto/gen/iop"
)

type noopHandler struct{}

func (h *noopHandler) OnRunRequest(_ context.Context, _ *transport.Session, _ *iop.RunRequest) error {
	return nil
}
func (h *noopHandler) OnCancel(_ context.Context, _ *transport.Session, _ *iop.CancelRequest) error {
	return nil
}
func (h *noopHandler) OnCommandRequest(_ context.Context, _ *transport.Session, _ *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
	return nil, nil
}
func (h *noopHandler) OnConfigRefresh(_ context.Context, _ *transport.Session, req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error) {
	return &iop.NodeConfigRefreshResponse{
		RequestId: req.GetRequestId(),
		Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED,
	}, nil
}
func (h *noopHandler) OnProviderTunnelRequest(_ context.Context, _ *transport.Session, _ *iop.ProviderTunnelRequest) error {
	return nil
}

func TestSession_SetHandler_ConcurrentSafe(t *testing.T) {
	var s transport.Session
	var wg sync.WaitGroup
	h := &noopHandler{}
	for i := 0; i < 50; i++ {
		wg.Add(1)
		go func() {
			t.Helper()
			defer wg.Done()
			s.SetHandler(h)
		}()
	}
	wg.Wait()
}

// TestSessionHealthObservationSeqIsMonotonicPerConnection verifies a new Session
// starts at zero, so the first finalized observation receives one and each
// subsequent call increments by one.
func TestSessionHealthObservationSeqIsMonotonicPerConnection(t *testing.T) {
	var s transport.Session
	for want := uint64(1); want <= 4; want++ {
		if got := s.NextHealthObservationSeq(); got != want {
			t.Fatalf("NextHealthObservationSeq() = %d, want %d", got, want)
		}
	}
}

// TestSessionHealthObservationSeqUniqueUnderConcurrency verifies concurrent
// normalized and tunnel attempts sharing one Session each receive a unique,
// contiguous value with no collisions or zeros.
func TestSessionHealthObservationSeqUniqueUnderConcurrency(t *testing.T) {
	var s transport.Session
	const workers = 64
	values := make(chan uint64, workers)
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			values <- s.NextHealthObservationSeq()
		}()
	}
	wg.Wait()
	close(values)
	seen := make(map[uint64]bool, workers)
	var maxSeq uint64
	for v := range values {
		if v == 0 {
			t.Fatal("finalized observation received sequence zero")
		}
		if seen[v] {
			t.Fatalf("health observation sequence %d issued twice", v)
		}
		seen[v] = true
		if v > maxSeq {
			maxSeq = v
		}
	}
	if len(seen) != workers || maxSeq != workers {
		t.Fatalf("concurrent sequence = %d distinct values, max %d; want %d contiguous", len(seen), maxSeq, workers)
	}
}

// TestSessionHealthObservationSeqResetsPerNewSession verifies the counter is
// connection-scoped: a second Session starts its own sequence at one regardless
// of how far the first advanced.
func TestSessionHealthObservationSeqResetsPerNewSession(t *testing.T) {
	var first, second transport.Session
	if got := first.NextHealthObservationSeq(); got != 1 {
		t.Fatalf("first session initial seq = %d, want 1", got)
	}
	first.NextHealthObservationSeq()
	first.NextHealthObservationSeq()
	if got := second.NextHealthObservationSeq(); got != 1 {
		t.Fatalf("second session initial seq = %d, want 1 (new connection starts at zero)", got)
	}
}

// TestSessionHealthObservationSeqWrapsMonotonically documents the overflow
// policy: the counter is monotonic within the uint64 space and wraps only after
// 2^64 observations on a single connection, which is unreachable in practice.
func TestSessionHealthObservationSeqWrapsMonotonically(t *testing.T) {
	var s transport.Session
	s.ExportSeedHealthObservationSeq(^uint64(0)) // 2^64 - 1
	if got := s.NextHealthObservationSeq(); got != 0 {
		t.Fatalf("wrap boundary seq = %d, want 0 after 2^64-1", got)
	}
	if got := s.NextHealthObservationSeq(); got != 1 {
		t.Fatalf("post-wrap seq = %d, want 1", got)
	}
}

// buildSessionTestPipe creates a net.Pipe-based pair: one side acts as "edge"
// (sends requests) and the other side acts as the node session under test.
// The edge side parser map must include the response type; the node side must
// include the request type (handled by nodeParserMap via DialEdge, but for
// unit tests we wire it manually).
func buildSessionTestPipe(t *testing.T) (edgeSide *toki.TcpClient, nodeSide *toki.TcpClient) {
	t.Helper()
	edgeConn, nodeConn := net.Pipe()
	edgeParserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunEvent{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunEvent{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeConfigRefreshResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeConfigRefreshResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.ProviderTunnelFrame{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelFrame{}
			return m, proto.Unmarshal(b, m)
		},
	}
	nodeParserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeConfigRefreshRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeConfigRefreshRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	edgeSide = toki.NewTcpClient(edgeConn, 0, 0, edgeParserMap)
	nodeSide = toki.NewTcpClient(nodeConn, 0, 0, nodeParserMap)
	t.Cleanup(func() { edgeSide.Close(); nodeSide.Close() })
	return edgeSide, nodeSide
}

// appliedHandler always returns applied.
type appliedHandler struct{ noopHandler }

func (h *appliedHandler) OnConfigRefresh(_ context.Context, _ *transport.Session, req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error) {
	return &iop.NodeConfigRefreshResponse{
		RequestId: req.GetRequestId(),
		Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED,
	}, nil
}

// errorHandler returns an error from OnConfigRefresh.
type errorHandler struct{ noopHandler }

func (h *errorHandler) OnConfigRefresh(_ context.Context, _ *transport.Session, req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error) {
	return nil, errors.New("refresh failed")
}

// TestSessionConfigRefreshRequestReturnsHandlerResponse verifies that a
// NodeConfigRefreshRequest pushed by the edge reaches the handler and the
// response is returned to the edge.
func TestSessionConfigRefreshRequestReturnsHandlerResponse(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)

	sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-test", "alias-test")
	sess.SetHandler(&appliedHandler{})

	resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
		&edgeSide.Communicator,
		&iop.NodeConfigRefreshRequest{RequestId: "req-1", ChangedPaths: []string{"nodes.0.providers.0.capacity"}},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped: %v", err)
	}
	if resp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
		t.Fatalf("expected status=applied, got %v", resp.GetStatus())
	}
	if resp.GetRequestId() != "req-1" {
		t.Fatalf("expected request_id=req-1, got %q", resp.GetRequestId())
	}
}

// TestSessionConfigRefreshRequestHandlerErrorReturnsFailure verifies that a
// handler error is translated to a failed protocol response.
func TestSessionConfigRefreshRequestHandlerErrorReturnsFailure(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)

	sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-test", "alias-test")
	sess.SetHandler(&errorHandler{})

	resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
		&edgeSide.Communicator,
		&iop.NodeConfigRefreshRequest{RequestId: "req-err"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped: %v", err)
	}
	if resp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED {
		t.Fatalf("expected status=failed, got %v", resp.GetStatus())
	}
	if resp.GetError() == "" {
		t.Fatal("expected non-empty error message")
	}
}

// TestSessionConfigRefreshRequestNoHandlerReturnsFailure verifies that when no
// handler is set, the session returns a failed response with an informative message.
func TestSessionConfigRefreshRequestNoHandlerReturnsFailure(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)

	_ = transport.ExportNewSession(nodeSide, zap.NewNop(), "node-nohandler", "")
	// handler intentionally NOT set

	resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
		&edgeSide.Communicator,
		&iop.NodeConfigRefreshRequest{RequestId: "req-nohandler"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped: %v", err)
	}
	if resp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED {
		t.Fatalf("expected status=failed when handler is nil, got %v", resp.GetStatus())
	}
}

type tunnelHandler struct {
	noopHandler
	mu       sync.Mutex
	requests []*iop.ProviderTunnelRequest
	done     chan struct{}
}

func (h *tunnelHandler) OnProviderTunnelRequest(_ context.Context, _ *transport.Session, req *iop.ProviderTunnelRequest) error {
	h.mu.Lock()
	h.requests = append(h.requests, req)
	h.mu.Unlock()
	close(h.done)
	return nil
}

func TestSessionProviderTunnelRequest(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)

	sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-test", "alias-test")
	handler := &tunnelHandler{done: make(chan struct{})}
	sess.SetHandler(handler)

	req := &iop.ProviderTunnelRequest{
		RunId:    "run-1",
		TunnelId: "tunnel-1",
		Adapter:  "openai_compat",
		Target:   "target-1",
	}

	if err := edgeSide.Send(req); err != nil {
		t.Fatalf("Send: %v", err)
	}

	select {
	case <-handler.done:
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for tunnel request")
	}

	handler.mu.Lock()
	defer handler.mu.Unlock()
	if len(handler.requests) != 1 {
		t.Fatalf("expected 1 request, got %d", len(handler.requests))
	}
	got := handler.requests[0]
	if got.GetRunId() != "run-1" || got.GetTunnelId() != "tunnel-1" {
		t.Errorf("unexpected request fields: %+v", got)
	}
}

type lifetimeHandler struct {
	noopHandler
	runStarted       chan struct{}
	runCanceled      chan error
	runSendResult    chan error
	tunnelStarted    chan struct{}
	tunnelCanceled   chan error
	tunnelSendResult chan error
}

func newLifetimeHandler() *lifetimeHandler {
	return &lifetimeHandler{
		runStarted: make(chan struct{}), runCanceled: make(chan error, 1), runSendResult: make(chan error, 1),
		tunnelStarted: make(chan struct{}), tunnelCanceled: make(chan error, 1), tunnelSendResult: make(chan error, 1),
	}
}

func (h *lifetimeHandler) OnRunRequest(ctx context.Context, sess *transport.Session, req *iop.RunRequest) error {
	close(h.runStarted)
	<-ctx.Done()
	h.runCanceled <- ctx.Err()
	h.runSendResult <- sess.Send(&iop.RunEvent{RunId: req.GetRunId(), Type: "error", Error: "must not reach dead session"})
	return ctx.Err()
}

func (h *lifetimeHandler) OnProviderTunnelRequest(ctx context.Context, sess *transport.Session, req *iop.ProviderTunnelRequest) error {
	close(h.tunnelStarted)
	<-ctx.Done()
	h.tunnelCanceled <- ctx.Err()
	h.tunnelSendResult <- sess.Send(&iop.ProviderTunnelFrame{RunId: req.GetRunId(), TunnelId: req.GetTunnelId(), Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR, Error: "must not reach dead session"})
	return ctx.Err()
}

func TestSessionLifetimeCancelsRunHandler(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)
	sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-run-lifetime", "alias")
	handler := newLifetimeHandler()
	sess.SetHandler(handler)
	if err := edgeSide.Send(&iop.RunRequest{RunId: "run-lifetime"}); err != nil {
		t.Fatal(err)
	}
	select {
	case <-handler.runStarted:
	case <-time.After(2 * time.Second):
		t.Fatal("run handler did not start")
	}
	if err := edgeSide.Close(); err != nil {
		t.Fatal(err)
	}
	select {
	case err := <-handler.runCanceled:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("run context error = %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("run handler context was not canceled on disconnect")
	}
	if err := <-handler.runSendResult; err == nil {
		t.Fatal("run terminal Send unexpectedly succeeded on dead session")
	}
}

func TestSessionLifetimeCancelsTunnelHandler(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)
	sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-tunnel-lifetime", "alias")
	handler := newLifetimeHandler()
	sess.SetHandler(handler)
	if err := edgeSide.Send(&iop.ProviderTunnelRequest{RunId: "run-tunnel-lifetime", TunnelId: "tunnel-lifetime"}); err != nil {
		t.Fatal(err)
	}
	select {
	case <-handler.tunnelStarted:
	case <-time.After(2 * time.Second):
		t.Fatal("tunnel handler did not start")
	}
	if err := edgeSide.Close(); err != nil {
		t.Fatal(err)
	}
	select {
	case err := <-handler.tunnelCanceled:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("tunnel context error = %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("tunnel handler context was not canceled on disconnect")
	}
	if err := <-handler.tunnelSendResult; err == nil {
		t.Fatal("tunnel terminal Send unexpectedly succeeded on dead session")
	}
}

// Compile check: Session must export a way to create instances for tests.
// ExportNewSession is expected in session_export_test.go or a separate test helper file.
var _ = fmt.Sprintf

type errorTunnelHandler struct {
	noopHandler
	done chan struct{}
}

func (h *errorTunnelHandler) OnProviderTunnelRequest(_ context.Context, _ *transport.Session, req *iop.ProviderTunnelRequest) error {
	defer close(h.done)
	return errors.New("tunnel error")
}

func TestSessionProviderTunnelRequest_NilAndErrHandler(t *testing.T) {
	// 1. Nil handler test
	{
		edgeSide, nodeSide := buildSessionTestPipe(t)
		sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-test-nil", "alias-test")
		// handler is nil

		req := &iop.ProviderTunnelRequest{
			RunId:    "run-nil",
			TunnelId: "tunnel-nil",
			Adapter:  "openai_compat",
			Target:   "target-nil",
		}

		if err := edgeSide.Send(req); err != nil {
			t.Fatalf("Send: %v", err)
		}

		// wait a bit to ensure no panic
		time.Sleep(100 * time.Millisecond)
		sess.Close()
	}

	// 2. Error handler test
	{
		edgeSide, nodeSide := buildSessionTestPipe(t)
		sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-test-err", "alias-test")
		handler := &errorTunnelHandler{done: make(chan struct{})}
		sess.SetHandler(handler)

		req := &iop.ProviderTunnelRequest{
			RunId:    "run-err",
			TunnelId: "tunnel-err",
			Adapter:  "openai_compat",
			Target:   "target-err",
		}

		if err := edgeSide.Send(req); err != nil {
			t.Fatalf("Send: %v", err)
		}

		select {
		case <-handler.done:
		case <-time.After(2 * time.Second):
			t.Fatal("timeout waiting for tunnel request")
		}
		sess.Close()
	}
}

// TestSessionSignalReady verifies the node→edge dispatch-ready handshake:
// SignalReady sends a NodeReadyRequest carrying the session's node id and returns
// nil only when the edge acks ready, and an error when the edge rejects the
// signal (stale connection) so the caller can tear down and reconnect.
func TestSessionSignalReady(t *testing.T) {
	newReadyPipe := func(t *testing.T, ready bool, reason string) *transport.Session {
		t.Helper()
		edgeConn, nodeConn := net.Pipe()
		edgeParser := toki.ParserMap{
			toki.TypeNameOf(&iop.NodeReadyRequest{}): func(b []byte) (proto.Message, error) {
				m := &iop.NodeReadyRequest{}
				return m, proto.Unmarshal(b, m)
			},
		}
		nodeParser := toki.ParserMap{
			toki.TypeNameOf(&iop.NodeReadyResponse{}): func(b []byte) (proto.Message, error) {
				m := &iop.NodeReadyResponse{}
				return m, proto.Unmarshal(b, m)
			},
		}
		edgeSide := toki.NewTcpClient(edgeConn, 0, 0, edgeParser)
		nodeSide := toki.NewTcpClient(nodeConn, 0, 0, nodeParser)
		t.Cleanup(func() { edgeSide.Close(); nodeSide.Close() })

		toki.AddRequestListenerTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
			&edgeSide.Communicator,
			func(req *iop.NodeReadyRequest) (*iop.NodeReadyResponse, error) {
				if req.GetNodeId() != "node-ready-test" {
					return &iop.NodeReadyResponse{Ready: false, Reason: "unexpected node id"}, nil
				}
				return &iop.NodeReadyResponse{Ready: ready, Reason: reason}, nil
			},
		)
		return transport.ExportNewSession(nodeSide, zap.NewNop(), "node-ready-test", "alias")
	}

	t.Run("ready ack succeeds", func(t *testing.T) {
		sess := newReadyPipe(t, true, "")
		if err := sess.SignalReady(2 * time.Second); err != nil {
			t.Fatalf("SignalReady on ready ack: %v", err)
		}
	})

	t.Run("non-ready ack errors", func(t *testing.T) {
		sess := newReadyPipe(t, false, "superseded")
		err := sess.SignalReady(2 * time.Second)
		if err == nil {
			t.Fatal("SignalReady must error when edge rejects the ready signal")
		}
	})
}