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

package service

import (
	"context"
	"fmt"
	"strings"
	"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
	ProviderID       string
	UsageAttribution string
	Adapter          string
	Target           string
	SessionID        string
	Method           string
	Path             string
	// Operation is the protocol operation id (e.g. "chat_completions",
	// "messages", "models"). When set, the Node adapter resolves the request
	// URL from the concrete profile's operation path. When empty, the legacy
	// Path field is used as a mixed-version fallback.
	Operation 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
	CredentialBinding    *CredentialBinding
}

// CredentialBinding contains only authenticated, secret-free route facts.
// The selected Node and target are added by Service after admission.
type CredentialBinding struct {
	PrincipalRef         string
	CredentialSlotRef    string
	RouteID              string
	ProfileID            string
	CredentialRevision   uint64
	RouteRevision        uint64
	ProjectionGeneration uint64
}

// 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{}
	}
	dispatch := h.RunDispatch
	dispatch.ProfileCapabilities = append([]string(nil), h.ProfileCapabilities...)
	return dispatch
}

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.submitProviderTunnelDirectContext(ctx, req)
}

func (s *Service) submitProviderTunnelQueued(ctx context.Context, req SubmitProviderTunnelRequest) (ProviderTunnelResult, error) {
	tunnelRunReq := SubmitRunRequest{
		NodeRef:              req.NodeRef,
		ModelGroupKey:        req.ModelGroupKey,
		ProviderID:           req.ProviderID,
		UsageAttribution:     req.UsageAttribution,
		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
	}
	if err := s.attachCredentialLease(ctx, req, selected.entry, target, tunnelReq); err != nil {
		reservation.release("credential-lease-error")
		return nil, err
	}
	if req.CredentialBinding != nil {
		defer s.releaseCredentialLease()
	}

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

	var handle *ProviderTunnelHandle
	err = s.registry.WithCurrentDispatchOwner(selected.entry.NodeID, selected.entry.Client, selected.generation, func() error {
		if err := s.validateCredentialFence(tunnelReq.GetCredentialBinding()); err != nil {
			return err
		}
		h, err := s.openProviderTunnel(selected.entry, tunnelReq, req, queueReason, true, selected.providerID, selected.providerType, string(selected.executionPath))
		if err != nil {
			return err
		}
		handle = h
		reservation.handOff()
		return nil
	})
	if err != nil {
		if !s.candidateIsCurrentOwner(selected) {
			reservation.release("stale-generation")
			return nil, staleGenerationError(selected)
		}
		reservation.release("send-error")
		return nil, err
	}
	return handle, nil
}

func (s *Service) submitProviderTunnelDirect(req SubmitProviderTunnelRequest) (ProviderTunnelResult, error) {
	return s.submitProviderTunnelDirectContext(context.Background(), req)
}

func (s *Service) submitProviderTunnelDirectContext(ctx context.Context, req SubmitProviderTunnelRequest) (ProviderTunnelResult, error) {
	entry, err := s.ResolveDispatchReady(req.NodeRef)
	if err != nil {
		return nil, err
	}
	tunnelReq, _, err := buildProviderTunnelRequest(req, req.Adapter, req.Target)
	if err != nil {
		return nil, err
	}
	if err := s.attachCredentialLease(ctx, req, entry, req.Target, tunnelReq); err != nil {
		return nil, err
	}
	if req.CredentialBinding != nil {
		defer s.releaseCredentialLease()
	}
	var handle *ProviderTunnelHandle
	err = s.registry.WithCurrentDispatchOwner(entry.NodeID, entry.Client, entry.ConnectionGeneration, func() error {
		if err := s.validateCredentialFence(tunnelReq.GetCredentialBinding()); err != nil {
			return err
		}
		var openErr error
		handle, openErr = s.openProviderTunnel(entry, tunnelReq, req, "dispatched", false, req.ProviderID, "", "")
		return openErr
	})
	if err != nil {
		return nil, err
	}
	return handle, nil
}

func (s *Service) attachCredentialLease(ctx context.Context, req SubmitProviderTunnelRequest, entry *edgenode.NodeEntry, target string, tunnelReq *iop.ProviderTunnelRequest) error {
	if req.CredentialBinding == nil {
		return nil
	}
	provider := s.credentialLeaseProvider()
	if provider == nil || entry == nil || entry.CredentialRecipientKeyID == "" || len(entry.CredentialRecipientPublicKey) != 32 {
		return fmt.Errorf("credential lease delivery unavailable")
	}
	if !s.reserveCredentialLease() {
		return fmt.Errorf("credential lease live set full")
	}
	reserved := true
	defer func() {
		if reserved {
			s.releaseCredentialLease()
		}
	}()
	binding := &iop.CredentialLeaseBinding{
		PrincipalRef: req.CredentialBinding.PrincipalRef, CredentialSlotRef: req.CredentialBinding.CredentialSlotRef,
		RouteId: req.CredentialBinding.RouteID, ProfileId: req.CredentialBinding.ProfileID,
		UpstreamTarget: target, NodeId: entry.NodeID, RecipientKeyId: entry.CredentialRecipientKeyID,
		CredentialRevision: req.CredentialBinding.CredentialRevision, RouteRevision: req.CredentialBinding.RouteRevision,
		ProjectionGeneration: req.CredentialBinding.ProjectionGeneration,
	}
	if err := provider.ValidateCredentialBinding(binding); err != nil {
		return fmt.Errorf("credential binding fenced")
	}
	lease, err := provider.AcquireCredentialLease(ctx, binding, entry.CredentialRecipientPublicKey)
	if err != nil {
		return err
	}
	if lease == nil || lease.GetScope() == nil {
		return fmt.Errorf("credential lease unavailable")
	}
	if !credentialLeaseScopeMatchesBinding(lease.GetScope(), binding) {
		return fmt.Errorf("credential lease scope mismatch")
	}
	for header := range tunnelReq.GetHeaders() {
		if strings.EqualFold(header, lease.GetScope().GetHeaderName()) {
			return fmt.Errorf("provider credential header collision")
		}
	}
	tunnelReq.CredentialBinding = binding
	tunnelReq.CredentialLease = lease
	reserved = false
	return nil
}

func credentialLeaseScopeMatchesBinding(scope *iop.CredentialLeaseScope, binding *iop.CredentialLeaseBinding) bool {
	if scope == nil || binding == nil {
		return false
	}
	return scope.GetPrincipalRef() == binding.GetPrincipalRef() &&
		scope.GetCredentialSlotRef() == binding.GetCredentialSlotRef() &&
		scope.GetRouteId() == binding.GetRouteId() &&
		scope.GetProfileId() == binding.GetProfileId() &&
		scope.GetUpstreamTarget() == binding.GetUpstreamTarget() &&
		scope.GetNodeId() == binding.GetNodeId() &&
		scope.GetRecipientKeyId() == binding.GetRecipientKeyId() &&
		scope.GetCredentialRevision() == binding.GetCredentialRevision() &&
		scope.GetRouteRevision() == binding.GetRouteRevision() &&
		scope.GetProjectionGeneration() == binding.GetProjectionGeneration()
}

func (s *Service) validateCredentialFence(binding *iop.CredentialLeaseBinding) error {
	if binding == nil {
		return nil
	}
	provider := s.credentialLeaseProvider()
	if provider == nil || provider.ValidateCredentialBinding(binding) != nil {
		return fmt.Errorf("credential binding fenced")
	}
	return nil
}

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

	credentialSlotRef := ""
	var credentialRevision uint64
	if lease := tunnelReq.GetCredentialLease(); lease != nil && lease.GetScope() != nil {
		credentialSlotRef = lease.GetScope().GetCredentialSlotRef()
		credentialRevision = lease.GetScope().GetCredentialRevision()
	}

	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,
			UsageAttribution:     req.UsageAttribution,
			ProviderType:         providerType,
			ExecutionPath:        executionPath,
			CredentialSlotRef:    credentialSlotRef,
			CredentialRevision:   credentialRevision,
			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,
		Operation:  req.Operation,
		Headers:    headers,
		Body:       body,
		Stream:     req.Stream,
		TimeoutSec: int32(normalizeTimeoutSec(req.TimeoutSec)),
		Metadata:   metadata,
		SessionId:  NormalizeSessionID(req.SessionID),
	}, runID, nil
}