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()
}

// 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.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.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)
	}
}

// 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")
		}
	})
}