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"
	"git.toki-labs.com/toki/proto-socket/go/packets"
	"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
}

// WorkspaceHandler is deliberately optional so existing provider Handler mocks
// and Node implementations remain source-compatible. The dedicated workspace
// boundary is not a RunRequest metadata extension or a NodeCommand variant.
type WorkspaceHandler interface {
	OnWorkspaceOpen(ctx context.Context, sess *Session, req *iop.WorkspaceOpenRequest) (*iop.WorkspaceOpenResponse, error)
	OnWorkspaceTool(ctx context.Context, sess *Session, req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error)
	OnWorkspaceCancel(ctx context.Context, sess *Session, req *iop.WorkspaceCancelRequest) (*iop.WorkspaceCancelResponse, error)
	OnWorkspaceCleanup(ctx context.Context, sess *Session, req *iop.WorkspaceCleanupRequest) (*iop.WorkspaceCleanupResponse, 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

	// workspaceResponseNonce sources the outgoing frame nonce for
	// asynchronously queued workspace responses. The peer matches replies purely
	// on the response nonce (the original request nonce), so this frame nonce is
	// informational; it stays a unique positive int32 per response to mirror the
	// communicator's own request/response framing.
	workspaceResponseNonce atomic.Int32
}

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

	s.registerWorkspaceListeners()
}

// registerWorkspaceListeners installs the four workspace request handlers. Unlike
// the shared AddRequestListenerTyped helper, which runs its callback synchronously
// on the communicator's single receive coordinator, each workspace request runs
// its handler and queues its typed response on a dedicated goroutine. Concurrency
// is required because a workspace tool handler may block until it observes its own
// cancellation, and a queued WorkspaceCancelRequest must still be dispatched while
// that tool handler is in flight. Request nonces, generic unsupported/failed
// responses, session-lifetime cancellation via s.Context(), and the optional
// WorkspaceHandler contract are all preserved.
func (s *Session) registerWorkspaceListeners() {
	addWorkspaceRequestListener(s, &iop.WorkspaceOpenRequest{}, func(req *iop.WorkspaceOpenRequest) proto.Message {
		workspace, ok := s.workspaceHandler()
		if !ok {
			return workspaceOpenUnsupported(req)
		}
		resp, err := workspace.OnWorkspaceOpen(s.Context(), s, req)
		if err != nil || resp == nil {
			return workspaceOpenFailed(req)
		}
		return resp
	})

	addWorkspaceRequestListener(s, &iop.WorkspaceToolRequest{}, func(req *iop.WorkspaceToolRequest) proto.Message {
		workspace, ok := s.workspaceHandler()
		if !ok {
			return workspaceToolUnsupported(req)
		}
		resp, err := workspace.OnWorkspaceTool(s.Context(), s, req)
		if err != nil || resp == nil {
			return workspaceToolFailed(req)
		}
		return resp
	})

	addWorkspaceRequestListener(s, &iop.WorkspaceCancelRequest{}, func(req *iop.WorkspaceCancelRequest) proto.Message {
		workspace, ok := s.workspaceHandler()
		if !ok {
			return workspaceCancelUnsupported(req)
		}
		resp, err := workspace.OnWorkspaceCancel(s.Context(), s, req)
		if err != nil || resp == nil {
			return workspaceCancelFailed(req)
		}
		return resp
	})

	addWorkspaceRequestListener(s, &iop.WorkspaceCleanupRequest{}, func(req *iop.WorkspaceCleanupRequest) proto.Message {
		workspace, ok := s.workspaceHandler()
		if !ok {
			return workspaceCleanupUnsupported(req)
		}
		resp, err := workspace.OnWorkspaceCleanup(s.Context(), s, req)
		if err != nil || resp == nil {
			return workspaceCleanupFailed(req)
		}
		return resp
	})
}

// addWorkspaceRequestListener registers a concurrent request-response handler for
// one workspace request type. The receive coordinator only routes the parsed
// request to a fresh goroutine, so a blocking handler never stalls other inbound
// frames on the connection. The queued response carries the original request
// nonce so the peer can match it; QueuePacket fails closed once the connection has
// drained, so a response produced after disconnect is dropped instead of written
// to a dead transport.
func addWorkspaceRequestListener[Req proto.Message](s *Session, sample Req, handle func(Req) proto.Message) {
	comm := &s.client.Communicator
	comm.AddRequestListener(toki.TypeNameOf(sample), func(m proto.Message, requestNonce int32) {
		req, ok := m.(Req)
		if !ok {
			return
		}
		go func() {
			resp := handle(req)
			data, err := proto.Marshal(resp)
			if err != nil {
				return
			}
			_ = comm.QueuePacket(&packets.PacketBase{
				TypeName:      toki.TypeNameOf(resp),
				Nonce:         s.nextWorkspaceResponseNonce(),
				ResponseNonce: requestNonce,
				Data:          data,
			})
		}()
	})
}

// nextWorkspaceResponseNonce returns a unique positive int32 for a workspace
// response frame. It wraps back to one on int32 overflow so the value stays
// positive like the communicator's own request nonces.
func (s *Session) nextWorkspaceResponseNonce() int32 {
	for {
		current := s.workspaceResponseNonce.Load()
		next := current + 1
		if next <= 0 {
			next = 1
		}
		if s.workspaceResponseNonce.CompareAndSwap(current, next) {
			return next
		}
	}
}

func (s *Session) workspaceHandler() (WorkspaceHandler, bool) {
	s.mu.RLock()
	handler := s.handler
	s.mu.RUnlock()
	workspace, ok := handler.(WorkspaceHandler)
	return workspace, ok && workspace != nil
}

func workspaceOpenUnsupported(req *iop.WorkspaceOpenRequest) *iop.WorkspaceOpenResponse {
	return &iop.WorkspaceOpenResponse{RequestId: req.GetRequestId(), WorkspaceRef: req.GetWorkspaceRef(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_UNSUPPORTED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_READY, Error: "workspace handler not ready"}
}

func workspaceOpenFailed(req *iop.WorkspaceOpenRequest) *iop.WorkspaceOpenResponse {
	return &iop.WorkspaceOpenResponse{RequestId: req.GetRequestId(), WorkspaceRef: req.GetWorkspaceRef(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, Error: "workspace handler failed"}
}

func workspaceToolUnsupported(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
	return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_UNSUPPORTED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_READY, Error: "workspace handler not ready"}
}

func workspaceToolFailed(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
	return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, Error: "workspace handler failed"}
}

func workspaceCancelUnsupported(req *iop.WorkspaceCancelRequest) *iop.WorkspaceCancelResponse {
	return &iop.WorkspaceCancelResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_UNSUPPORTED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_READY, Error: "workspace handler not ready"}
}

func workspaceCancelFailed(req *iop.WorkspaceCancelRequest) *iop.WorkspaceCancelResponse {
	return &iop.WorkspaceCancelResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, Error: "workspace handler failed"}
}

func workspaceCleanupUnsupported(req *iop.WorkspaceCleanupRequest) *iop.WorkspaceCleanupResponse {
	return &iop.WorkspaceCleanupResponse{RequestId: req.GetRequestId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_UNSUPPORTED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_READY, Error: "workspace handler not ready"}
}

func workspaceCleanupFailed(req *iop.WorkspaceCleanupRequest) *iop.WorkspaceCleanupResponse {
	return &iop.WorkspaceCleanupResponse{RequestId: req.GetRequestId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, Error: "workspace handler failed"}
}

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