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

package transport

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

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

	"iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

// Handler processes IOP messages received from edge.
type Handler interface {
	OnRunRequest(ctx context.Context, sess *Session, req *iop.RunRequest) error
	OnCancel(ctx context.Context, sess *Session, req *iop.CancelRequest) error
	OnCommandRequest(ctx context.Context, sess *Session, req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error)
	OnConfigRefresh(ctx context.Context, sess *Session, req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error)
	OnProviderTunnelRequest(ctx context.Context, sess *Session, req *iop.ProviderTunnelRequest) error
}

// Session represents the node's persistent connection to edge.
type Session struct {
	client         *toki.TcpClient
	logger         *zap.Logger
	nodeID         string
	alias          string
	mu             sync.RWMutex
	handler        Handler
	eventHandler   func(*iop.EdgeNodeEvent)
	closeReason    string
	disconnectCh   chan struct{}
	disconnectOnce sync.Once
	lifetimeCtx    context.Context
	lifetimeCancel context.CancelFunc

	// healthObservationSeq is the connection-scoped source of monotonic
	// health-observation sequence numbers. A new Session starts at zero, so the
	// first finalized observation receives one. Normalized and tunnel attempts
	// on the same Session share this source and receive unique, monotonically
	// increasing values under concurrency. It never resets within a connection
	// and never encodes a process-global generation.
	healthObservationSeq atomic.Uint64
}

func newSession(client *toki.TcpClient, logger *zap.Logger, nodeID, alias string) *Session {
	lifetimeCtx, lifetimeCancel := context.WithCancel(context.Background())
	s := &Session{client: client, logger: logger, nodeID: nodeID, alias: alias, disconnectCh: make(chan struct{}), lifetimeCtx: lifetimeCtx, lifetimeCancel: lifetimeCancel}
	s.registerExecutionListeners()
	s.registerControlListeners()
	s.registerConnectionListeners()
	return s
}

func (s *Session) registerExecutionListeners() {
	toki.AddListenerTyped[*iop.RunRequest](&s.client.Communicator, func(req *iop.RunRequest) {
		go func() {
			s.mu.RLock()
			h := s.handler
			s.mu.RUnlock()
			if h == nil {
				return
			}
			if err := h.OnRunRequest(s.Context(), s, req); err != nil {
				s.logger.Warn("run request error",
					zap.String("run_id", req.GetRunId()),
					zap.Error(err),
				)
			}
		}()
	})

	toki.AddListenerTyped[*iop.ProviderTunnelRequest](&s.client.Communicator, func(req *iop.ProviderTunnelRequest) {
		go func() {
			s.mu.RLock()
			h := s.handler
			s.mu.RUnlock()
			if h == nil {
				s.logger.Warn("provider tunnel request ignored: handler not ready",
					zap.String("run_id", req.GetRunId()),
				)
				return
			}
			if err := h.OnProviderTunnelRequest(s.Context(), s, req); err != nil {
				s.logger.Warn("provider tunnel request error",
					zap.String("run_id", req.GetRunId()),
					zap.Error(err),
				)
			}
		}()
	})
}

func (s *Session) registerControlListeners() {
	toki.AddListenerTyped[*iop.CancelRequest](&s.client.Communicator, func(req *iop.CancelRequest) {
		s.mu.RLock()
		h := s.handler
		s.mu.RUnlock()
		if h == nil {
			return
		}
		if err := h.OnCancel(context.Background(), s, req); err != nil {
			s.logger.Warn("cancel error", zap.String("run_id", req.GetRunId()), zap.Error(err))
		}
	})

	toki.AddRequestListenerTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](&s.client.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
		s.mu.RLock()
		h := s.handler
		s.mu.RUnlock()
		if h == nil {
			return &iop.NodeCommandResponse{Error: "handler not ready"}, nil
		}
		resp, err := h.OnCommandRequest(context.Background(), s, req)
		if err != nil {
			return &iop.NodeCommandResponse{Error: err.Error()}, nil
		}
		return resp, nil
	})

	toki.AddRequestListenerTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](&s.client.Communicator, func(req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error) {
		s.mu.RLock()
		h := s.handler
		s.mu.RUnlock()
		if h == nil {
			return &iop.NodeConfigRefreshResponse{
				RequestId: req.GetRequestId(),
				Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED,
				Error:     "handler not ready",
			}, nil
		}
		resp, err := h.OnConfigRefresh(context.Background(), s, req)
		if err != nil {
			return &iop.NodeConfigRefreshResponse{
				RequestId: req.GetRequestId(),
				Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED,
				Error:     err.Error(),
			}, nil
		}
		return resp, nil
	})
}

func (s *Session) registerConnectionListeners() {
	toki.AddListenerTyped[*iop.EdgeNodeEvent](&s.client.Communicator, func(event *iop.EdgeNodeEvent) {
		s.emitEvent(event)
	})

	s.client.AddDisconnectListener(func(_ *toki.TcpClient) {
		transportInfo := s.client.DisconnectInfo()
		s.logger.Info("disconnected from edge", transportDisconnectFields(transportInfo)...)
		s.emitEvent(events.NewEdgeNodeEvent(
			events.SourceNode,
			events.TypeEdgeDisconnected,
			s.nodeID,
			s.alias,
			s.disconnectReason(),
			transportDisconnectMetadata(transportInfo),
		))
		s.disconnectOnce.Do(func() { s.lifetimeCancel(); close(s.disconnectCh) })
	})
}

// SetHandler attaches the message handler. Called after registration completes.
func (s *Session) SetHandler(h Handler) {
	s.mu.Lock()
	s.handler = h
	s.mu.Unlock()
}

// SignalReady tells edge this node has applied the config from RegisterResponse
// and installed its message handler, so edge may now open dispatch eligibility
// and pump any waiters stranded while the node was offline. It MUST be called
// after SetHandler: edge dispatches run/tunnel requests in response to this
// signal, and a request arriving before the handler is installed would be
// dropped. A non-ready ack (stale connection) or a transport error is returned so
// the caller tears the session down and lets the supervisor reconnect.
func (s *Session) SignalReady(timeout time.Duration) error {
	resp, err := toki.SendRequestTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
		&s.client.Communicator,
		&iop.NodeReadyRequest{NodeId: s.nodeID},
		timeout,
	)
	if err != nil {
		return err
	}
	if !resp.GetReady() {
		return fmt.Errorf("edge rejected ready signal: %s", resp.GetReason())
	}
	return nil
}

// NextHealthObservationSeq allocates the next connection-scoped health
// observation sequence value. It is atomic, so concurrent normalized and tunnel
// attempts on the same Session each receive a unique, monotonically increasing
// value; a new Session starts at zero, so the first observation receives one.
// The counter is monotonic within the uint64 space and wraps only after 2^64
// observations on a single connection, which is unreachable in practice. It is
// evidence sequencing only and never advances original request progress.
func (s *Session) NextHealthObservationSeq() uint64 {
	return s.healthObservationSeq.Add(1)
}

// NodeID returns the session's node ID.
func (s *Session) NodeID() string {
	return s.nodeID
}

// Alias returns the session's node alias.
func (s *Session) Alias() string {
	return s.alias
}

func (s *Session) SetEventHandler(handler func(*iop.EdgeNodeEvent)) {
	s.mu.Lock()
	s.eventHandler = handler
	s.mu.Unlock()
}

// Send transmits a proto message to edge.
func (s *Session) Send(m proto.Message) error {
	return s.client.Send(m)
}

// IsAlive reports whether the connection is active.
func (s *Session) IsAlive() bool {
	if s == nil || s.client == nil {
		return false
	}
	return s.client.IsAlive()
}

// Done returns a channel that is closed when the session disconnects (local or remote).
func (s *Session) Done() <-chan struct{} {
	if s == nil || s.disconnectCh == nil {
		return nil
	}
	return s.disconnectCh
}

// Context is canceled exactly once when this connection closes. Request
// handlers derive their per-request context from it, so a dead connection
// cannot retain an active provider attempt.
func (s *Session) Context() context.Context {
	if s == nil || s.lifetimeCtx == nil {
		return context.Background()
	}
	return s.lifetimeCtx
}

// IsLocalShutdown reports whether the disconnect was initiated by a local Close call.
func (s *Session) IsLocalShutdown() bool {
	return s.disconnectReason() == events.ReasonLocalShutdown
}

// Close terminates the connection to edge.
func (s *Session) Close() error {
	s.setCloseReason(events.ReasonLocalShutdown)
	return s.client.Close()
}

func (s *Session) emitEvent(event *iop.EdgeNodeEvent) {
	s.mu.RLock()
	handler := s.eventHandler
	s.mu.RUnlock()
	if handler != nil {
		handler(event)
	}
}

func (s *Session) setCloseReason(reason string) {
	s.mu.Lock()
	if s.closeReason == "" {
		s.closeReason = reason
	}
	s.mu.Unlock()
}

func (s *Session) disconnectReason() string {
	s.mu.RLock()
	reason := s.closeReason
	s.mu.RUnlock()
	if reason == "" {
		return events.ReasonTransportClosed
	}
	return reason
}

func transportDisconnectMetadata(info toki.DisconnectInfo) map[string]string {
	metadata := make(map[string]string, 2)
	if info.Reason != "" {
		metadata[events.MetadataTransportCloseReason] = info.Reason
	}
	if info.Error != "" {
		metadata[events.MetadataTransportCloseError] = info.Error
	}
	if len(metadata) == 0 {
		return nil
	}
	return metadata
}

func transportDisconnectFields(info toki.DisconnectInfo) []zap.Field {
	fields := make([]zap.Field, 0, 2)
	if info.Reason != "" {
		fields = append(fields, zap.String("transport_close_reason", info.Reason))
	}
	if info.Error != "" {
		fields = append(fields, zap.String("transport_close_error", info.Error))
	}
	return fields
}

// ExportNewSession exposes newSession for black-box transport and node tests.
func ExportNewSession(client *toki.TcpClient, logger *zap.Logger, nodeID, alias string) *Session {
	return newSession(client, logger, nodeID, alias)
}

// ExportSeedHealthObservationSeq presets the connection-scoped health
// observation counter for black-box tests that must exercise the monotonic wrap
// boundary without allocating 2^64 values.
func (s *Session) ExportSeedHealthObservationSeq(value uint64) {
	s.healthObservationSeq.Store(value)
}