iop/apps/edge/internal/service/provider_tunnel.go

package service

import (
	"context"
	"sync"
	"time"

	edgenode "iop/apps/edge/internal/node"
	iop "iop/proto/gen/iop"
)

// tunnelFrameBuffer is the per-tunnel channel capacity. The router blocks on a
// full buffer instead of dropping so ordered passthrough bytes are never lost
// while the subscriber is alive.
const tunnelFrameBuffer = 4096

// providerTunnelRouter delivers ProviderTunnelFrame messages to the
// request-bound subscriber of their tunnel id. It is intentionally separate
// from events.Bus: tunnel body bytes are the passthrough source of truth and
// must not enter the lossy run event fanout.
type providerTunnelRouter struct {
	mu   sync.Mutex
	subs map[string]*providerTunnelSub
}

type providerTunnelSub struct {
	ch   chan *iop.ProviderTunnelFrame
	done chan struct{}
}

func newProviderTunnelRouter() *providerTunnelRouter {
	return &providerTunnelRouter{subs: make(map[string]*providerTunnelSub)}
}

func (r *providerTunnelRouter) subscribe(tunnelID string, buffer int) (<-chan *iop.ProviderTunnelFrame, func()) {
	sub := &providerTunnelSub{
		ch:   make(chan *iop.ProviderTunnelFrame, buffer),
		done: make(chan struct{}),
	}
	r.mu.Lock()
	r.subs[tunnelID] = sub
	r.mu.Unlock()
	var once sync.Once
	unsubscribe := func() {
		once.Do(func() {
			r.mu.Lock()
			if cur, ok := r.subs[tunnelID]; ok && cur == sub {
				delete(r.subs, tunnelID)
			}
			r.mu.Unlock()
			close(sub.done)
		})
	}
	return sub.ch, unsubscribe
}

// route delivers a frame to its tunnel subscriber. Frames without a subscriber
// are dropped (the request already finished or was cancelled).
func (r *providerTunnelRouter) route(frame *iop.ProviderTunnelFrame) {
	if frame == nil {
		return
	}
	r.mu.Lock()
	sub := r.subs[frame.GetTunnelId()]
	r.mu.Unlock()
	if sub == nil {
		return
	}
	select {
	case sub.ch <- frame:
	case <-sub.done:
	}
}

// RouteProviderTunnelFrame delivers a ProviderTunnelFrame received from Edge
// transport to its request-bound tunnel stream. It never publishes to the
// run/node event bus.
func (s *Service) RouteProviderTunnelFrame(frame *iop.ProviderTunnelFrame) {
	if s == nil || s.tunnels == nil {
		return
	}
	s.tunnels.route(frame)
}

// SubmitProviderTunnelRequest asks a node to open a raw provider HTTP request
// and relay the response as ordered ProviderTunnelFrame messages. It is the
// passthrough sibling of SubmitRunRequest and shares the provider-pool
// admission gate.
type SubmitProviderTunnelRequest struct {
	NodeRef       string
	RunID         string
	ModelGroupKey string
	Adapter       string
	Target        string
	SessionID     string
	Method        string
	Path          string
	Headers       map[string]string
	Body          []byte
	// BuildBody, when set, produces the provider request body from the final
	// resolved target (provider-pool admission rewrites the target to the
	// winning candidate's served model). It takes precedence over Body.
	BuildBody            func(target string) ([]byte, error)
	Stream               bool
	TimeoutSec           int
	MaxQueue             int
	QueueTimeoutMS       int
	Metadata             map[string]string
	EstimatedInputTokens int
	ContextClass         string
	ProviderPool         bool
}

// ProviderTunnelStream carries the ordered raw provider frames of a dispatched
// tunnel. The channel is closed after the terminal END/ERROR frame or Close.
type ProviderTunnelStream struct {
	Frames <-chan *iop.ProviderTunnelFrame
}

// ProviderTunnelResult is the surface-neutral handle for a dispatched provider
// tunnel, mirroring RunResult for the raw passthrough path.
type ProviderTunnelResult interface {
	Dispatch() RunDispatch
	Stream() ProviderTunnelStream
	Close()
	WaitTimeout() time.Duration
	// SetHeaders allows late-binding of tunnel headers (e.g. provider auth
	// validated after provider-pool selection). Implementations may ignore
	// this if headers are already baked into the request.
	SetHeaders(map[string]string)
}

// ProviderTunnelHandle implements ProviderTunnelResult for tunnels dispatched
// over the Edge-Node socket.
type ProviderTunnelHandle struct {
	RunDispatch
	TunnelID  string
	Headers   map[string]string
	frames    <-chan *iop.ProviderTunnelFrame
	closeOnce sync.Once
	close     func()
}

func (h *ProviderTunnelHandle) Dispatch() RunDispatch {
	if h == nil {
		return RunDispatch{}
	}
	return h.RunDispatch
}

func (h *ProviderTunnelHandle) Stream() ProviderTunnelStream {
	if h == nil {
		return ProviderTunnelStream{}
	}
	return ProviderTunnelStream{Frames: h.frames}
}

func (h *ProviderTunnelHandle) Close() {
	if h != nil && h.close != nil {
		h.closeOnce.Do(h.close)
	}
}

func (h *ProviderTunnelHandle) WaitTimeout() time.Duration {
	if h == nil {
		return time.Duration(DefaultTimeoutSec+5) * time.Second
	}
	return time.Duration(normalizeTimeoutSec(h.TimeoutSec)+5) * time.Second
}

func (h *ProviderTunnelHandle) SetHeaders(hdrs map[string]string) {
	h.Headers = hdrs
}

// SubmitProviderTunnel dispatches a raw provider tunnel request. Provider-pool
// requests go through the same admission gate as SubmitRun; the reserved slot
// is released when the tunnel reaches END/ERROR or the handle is closed
// (cancel), never via the run event bus.
func (s *Service) SubmitProviderTunnel(ctx context.Context, req SubmitProviderTunnelRequest) (ProviderTunnelResult, error) {
	if req.ProviderPool && req.ModelGroupKey != "" && s.queue != nil {
		return s.submitProviderTunnelQueued(ctx, req)
	}
	return s.submitProviderTunnelDirect(req)
}

func (s *Service) submitProviderTunnelQueued(ctx context.Context, req SubmitProviderTunnelRequest) (ProviderTunnelResult, error) {
	tunnelRunReq := SubmitRunRequest{
		NodeRef:              req.NodeRef,
		ModelGroupKey:        req.ModelGroupKey,
		Adapter:              req.Adapter,
		Target:               req.Target,
		MaxQueue:             req.MaxQueue,
		QueueTimeoutMS:       req.QueueTimeoutMS,
		EstimatedInputTokens: req.EstimatedInputTokens,
		ContextClass:         req.ContextClass,
		ProviderPool:         true,
	}
	candidates, returnedPolicy, err := s.resolveQueueCandidates(tunnelRunReq)
	if err != nil {
		return nil, err
	}

	long := req.ContextClass == contextClassLong

	// Provider-pool tunnel uses the canonical policy from the runtime snapshot.
	var policy groupPolicy
	if tunnelRunReq.ProviderPool {
		_, _, policy = s.runtimeConfigSnapshot()
	} else {
		policy = returnedPolicy
	}

	selected, queueReason, err := s.queue.admitWithReason(ctx, req.ModelGroupKey, req.Adapter, req.Target, candidates, policy, s.resolveQueueCandidatesClosure(tunnelRunReq), long, tunnelRunReq.ProviderPool)
	if err != nil {
		return nil, err
	}
	// The admitted slot is owned by one reservation from here on: every failure
	// path below releases through it, and a dispatched tunnel hands it off to
	// its terminal frame / close path.
	reservation := newQueueReservation(s.queue, selected)

	adapter := req.Adapter
	if selected.adapter != "" {
		adapter = selected.adapter
	}
	target := req.Target
	if selected.servedTarget != "" {
		target = selected.servedTarget
	}

	tunnelReq, runID, err := buildProviderTunnelRequest(req, adapter, target)
	if err != nil {
		reservation.release("build-error")
		return nil, err
	}

	// Track inflight before send so END/ERROR/close release paths can find it
	// even if the terminal frame arrives before Send returns.
	reservation.track(runID)

	handle, err := s.openProviderTunnel(selected.entry, tunnelReq, req, queueReason, true, selected.providerID, selected.providerType, string(selected.executionPath))
	if err != nil {
		reservation.release("send-error")
		return nil, err
	}
	reservation.handOff()
	return handle, nil
}

func (s *Service) submitProviderTunnelDirect(req SubmitProviderTunnelRequest) (ProviderTunnelResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return nil, err
	}
	tunnelReq, _, err := buildProviderTunnelRequest(req, req.Adapter, req.Target)
	if err != nil {
		return nil, err
	}
	return s.openProviderTunnel(entry, tunnelReq, req, "dispatched", false, "", "", "")
}

// openProviderTunnel subscribes the request-bound frame channel, sends the
// tunnel request to the node, and wraps the raw channel so the reserved
// provider-pool slot is released exactly once on END/ERROR/close.
func (s *Service) openProviderTunnel(entry *edgenode.NodeEntry, tunnelReq *iop.ProviderTunnelRequest, req SubmitProviderTunnelRequest, queueReason string, tracked bool, providerID, providerType, executionPath string) (*ProviderTunnelHandle, error) {
	raw, unsubscribe := s.tunnels.subscribe(tunnelReq.GetTunnelId(), tunnelFrameBuffer)
	if err := entry.Client.Send(tunnelReq); err != nil {
		unsubscribe()
		return nil, err
	}

	runID := tunnelReq.GetRunId()
	var releaseOnce sync.Once
	release := func(reason string) {
		releaseOnce.Do(func() {
			if tracked && s.queue != nil {
				s.queue.releaseRun(runID, reason)
			}
		})
	}

	out := make(chan *iop.ProviderTunnelFrame, tunnelFrameBuffer)
	done := make(chan struct{})
	go func() {
		defer close(out)
		for {
			select {
			case frame, ok := <-raw:
				if !ok {
					release("tunnel-stream-closed")
					return
				}
				select {
				case out <- frame:
				case <-done:
					release("tunnel-closed")
					return
				}
				if isTerminalProviderTunnelFrame(frame) {
					release("tunnel-" + frame.GetKind().String())
					return
				}
			case <-done:
				release("tunnel-closed")
				return
			}
		}
	}()

	return &ProviderTunnelHandle{
		RunDispatch: RunDispatch{
			RunID:                runID,
			NodeID:               entry.NodeID,
			NodeLabel:            nodeLabel(entry),
			ModelGroupKey:        req.ModelGroupKey,
			Adapter:              tunnelReq.GetAdapter(),
			Target:               tunnelReq.GetTarget(),
			SessionID:            tunnelReq.GetSessionId(),
			TimeoutSec:           int(tunnelReq.GetTimeoutSec()),
			EstimatedInputTokens: req.EstimatedInputTokens,
			ContextClass:         req.ContextClass,
			ProviderID:           providerID,
			ProviderType:         providerType,
			ExecutionPath:        executionPath,
			QueueReason:          queueReason,
		},
		TunnelID: tunnelReq.GetTunnelId(),
		frames:   out,
		close: func() {
			close(done)
			unsubscribe()
			release("tunnel-closed")
		},
	}, nil
}

func isTerminalProviderTunnelFrame(f *iop.ProviderTunnelFrame) bool {
	switch f.GetKind() {
	case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
		iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
		return true
	default:
		return false
	}
}

func buildProviderTunnelRequest(req SubmitProviderTunnelRequest, adapter, target string) (*iop.ProviderTunnelRequest, string, error) {
	runID := req.RunID
	if runID == "" {
		runID = NewRunID()
	}
	body := req.Body
	if req.BuildBody != nil {
		b, err := req.BuildBody(target)
		if err != nil {
			return nil, "", err
		}
		body = b
	}
	headers := make(map[string]string, len(req.Headers))
	for k, v := range req.Headers {
		headers[k] = v
	}
	metadata := make(map[string]string, len(req.Metadata))
	for k, v := range req.Metadata {
		metadata[k] = v
	}
	return &iop.ProviderTunnelRequest{
		RunId:      runID,
		TunnelId:   runID + "-tunnel",
		Adapter:    adapter,
		Target:     target,
		Method:     req.Method,
		Path:       req.Path,
		Headers:    headers,
		Body:       body,
		Stream:     req.Stream,
		TimeoutSec: int32(normalizeTimeoutSec(req.TimeoutSec)),
		Metadata:   metadata,
		SessionId:  NormalizeSessionID(req.SessionID),
	}, runID, nil
}