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

package transport

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

	"go.uber.org/zap"

	"iop/packages/protocol"
)

const (
	heartbeatInterval = 30 * time.Second
	heartbeatTimeout  = 10 * time.Second
)

// Handler handles decoded messages arriving on a Session.
type Handler interface {
	OnRunRequest(ctx context.Context, sess *Session, req *protocol.RunRequest) error
	OnCapabilityRequest(ctx context.Context, sess *Session) (*protocol.CapabilityResponse, error)
	OnCancel(ctx context.Context, sess *Session, req *protocol.CancelRequest) error
}

// Session represents one active TCP connection from a client.
type Session struct {
	id      string
	conn    net.Conn
	handler Handler
	logger  *zap.Logger

	writeMu sync.Mutex
	closed  atomic.Bool

	cancelFns sync.Map // run_id → context.CancelFunc
}

func newSession(id string, conn net.Conn, handler Handler, logger *zap.Logger) *Session {
	return &Session{
		id:      id,
		conn:    conn,
		handler: handler,
		logger:  logger.With(zap.String("session_id", id)),
	}
}

// ID returns the session identifier.
func (s *Session) ID() string { return s.id }

// Send encodes env as a frame and writes it to the connection.
func (s *Session) Send(env *protocol.Envelope) error {
	if s.closed.Load() {
		return fmt.Errorf("session %s: closed", s.id)
	}
	payload, err := encodeEnvelope(env)
	if err != nil {
		return err
	}
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	return writeFrame(s.conn, payload)
}

// Close terminates the session.
func (s *Session) Close() {
	if s.closed.CompareAndSwap(false, true) {
		s.conn.Close()
		s.logger.Info("session closed")
	}
}

// RegisterCancel stores a cancel function for a run, enabling CancelRequest handling.
func (s *Session) RegisterCancel(runID string, cancel context.CancelFunc) {
	s.cancelFns.Store(runID, cancel)
}

// DeregisterCancel removes a run's cancel function after completion.
func (s *Session) DeregisterCancel(runID string) {
	s.cancelFns.Delete(runID)
}

// CancelRun invokes the cancel function for a run if one is registered.
func (s *Session) CancelRun(runID string) {
	if v, ok := s.cancelFns.Load(runID); ok {
		v.(context.CancelFunc)()
	}
}

// readLoop drives the session until the connection closes or an error occurs.
func (s *Session) readLoop(ctx context.Context) {
	defer s.Close()

	// Perform capability handshake: server sends capability request,
	// client responds (optional — here we accept any client proactively).
	go s.heartbeatLoop(ctx)

	for {
		if s.closed.Load() {
			return
		}

		raw, err := readFrame(s.conn)
		if err != nil {
			if !s.closed.Load() {
				s.logger.Debug("read frame error", zap.Error(err))
			}
			return
		}
		if len(raw) == 0 {
			continue
		}

		env, err := decodeEnvelope(raw)
		if err != nil {
			s.logger.Warn("decode envelope failed", zap.Error(err))
			continue
		}

		if err := s.dispatch(ctx, env); err != nil {
			s.logger.Warn("dispatch error", zap.String("type", string(env.Type)), zap.Error(err))
		}
	}
}

func (s *Session) dispatch(ctx context.Context, env *protocol.Envelope) error {
	switch env.Type {
	case protocol.TypeHeartbeat:
		return s.handleHeartbeat(env)

	case protocol.TypeRunRequest:
		var req protocol.RunRequest
		if err := decodePayload(env.Payload, &req); err != nil {
			return s.sendError(env.RequestID, "DECODE_ERROR", err.Error())
		}
		go func() {
			if err := s.handler.OnRunRequest(ctx, s, &req); err != nil {
				_ = s.sendError(env.RequestID, "RUN_ERROR", err.Error())
			}
		}()

	case protocol.TypeCapabilityRequest:
		resp, err := s.handler.OnCapabilityRequest(ctx, s)
		if err != nil {
			return s.sendError(env.RequestID, "CAPABILITY_ERROR", err.Error())
		}
		return s.sendCapabilityResponse(env.RequestID, resp)

	case protocol.TypeCancelRequest:
		var req protocol.CancelRequest
		if err := decodePayload(env.Payload, &req); err != nil {
			return s.sendError(env.RequestID, "DECODE_ERROR", err.Error())
		}
		return s.handler.OnCancel(ctx, s, &req)

	default:
		s.logger.Warn("unknown message type", zap.String("type", string(env.Type)))
	}
	return nil
}

func (s *Session) handleHeartbeat(env *protocol.Envelope) error {
	// Echo heartbeat back
	resp, err := encodePayload(protocol.Heartbeat{Timestamp: time.Now().UnixNano()})
	if err != nil {
		return err
	}
	return s.Send(&protocol.Envelope{
		ProtocolVersion: protocol.ProtocolVersion,
		RequestID:       env.RequestID,
		Type:            protocol.TypeHeartbeat,
		Payload:         resp,
	})
}

func (s *Session) sendError(requestID, code, message string) error {
	payload, err := encodePayload(protocol.ErrorMsg{Code: code, Message: message})
	if err != nil {
		return err
	}
	return s.Send(&protocol.Envelope{
		ProtocolVersion: protocol.ProtocolVersion,
		RequestID:       requestID,
		Type:            protocol.TypeError,
		Payload:         payload,
	})
}

func (s *Session) sendCapabilityResponse(requestID string, resp *protocol.CapabilityResponse) error {
	payload, err := encodePayload(resp)
	if err != nil {
		return err
	}
	return s.Send(&protocol.Envelope{
		ProtocolVersion: protocol.ProtocolVersion,
		RequestID:       requestID,
		Type:            protocol.TypeCapabilityResponse,
		Payload:         payload,
	})
}

// heartbeatLoop periodically sends heartbeat frames to detect dead connections.
func (s *Session) heartbeatLoop(ctx context.Context) {
	ticker := time.NewTicker(heartbeatInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if s.closed.Load() {
				return
			}
			payload, _ := encodePayload(protocol.Heartbeat{Timestamp: time.Now().UnixNano()})
			if err := s.Send(&protocol.Envelope{
				ProtocolVersion: protocol.ProtocolVersion,
				Type:            protocol.TypeHeartbeat,
				Payload:         payload,
			}); err != nil {
				s.logger.Debug("heartbeat send failed", zap.Error(err))
				s.Close()
				return
			}
		}
	}
}