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

package service

import (
	"context"
	"fmt"
)

// providerPoolPath indicates which execution path was selected for a
// provider-pool dispatch. These constants mirror the candidateNode.executionPath
// values but are exposed at the dispatch surface so callers (e.g. OpenAI
// handler) can distinguish tunnel/passthrough from normalized execution.
type providerPoolPath string

const (
	// ProviderPoolPathTunnel means the selected candidate executes via
	// provider tunnel / OpenAI-compatible passthrough.
	ProviderPoolPathTunnel providerPoolPath = "provider_tunnel"
	// ProviderPoolPathNormalized means the selected candidate executes via
	// normalized RunEvent path (Ollama/CLI/native).
	ProviderPoolPathNormalized providerPoolPath = "normalized"
)

// PrepareTunnel is an optional pre-dispatch hook that lets the caller
// inject or modify headers on a tunnel-path request before buildProviderTunnelRequest
// and the Node Send step run. If PrepareTunnel returns an error, the slot is
// released and no tunnel request is sent. This avoids late header mutation
// after wire dispatch for provider-pool tunnel paths.
type prepareTunnelFunc func(req SubmitProviderTunnelRequest) (SubmitProviderTunnelRequest, error)

// PrepareRun is an optional pre-dispatch hook that lets the caller prepare
// or validate the Run request for the normalized execution path. It runs
// AFTER a candidate is selected and ONLY on the normalized path (not on
// tunnel/passthrough). If PrepareRun returns an error, the slot is released
// and no normalized run is sent. This lets the caller enforce that a complete
// SubmitRunRequest (including ModelGroupKey and other required fields) is
// present before dispatch, while leaving the tunnel branch untouched.
type prepareRunFunc func(req SubmitRunRequest) (SubmitRunRequest, error)

// ProviderPoolDispatchRequest bundles the Run and Tunnel surface values for
// a single one-shot provider-pool dispatch. SubmitProviderPool uses exactly
// one queue admission to select a candidate, then dispatches only the
// execution path indicated by the candidate's executionPath.
type ProviderPoolDispatchRequest struct {
	Run           SubmitRunRequest
	Tunnel        SubmitProviderTunnelRequest
	PrepareTunnel prepareTunnelFunc
	PrepareRun    prepareRunFunc
}

// ProviderPoolDispatchResult describes which execution path was selected and
// carries the corresponding dispatch result. Exactly one of Run or Tunnel is
// non-nil.
type ProviderPoolDispatchResult struct {
	Path         providerPoolPath
	Run          RunResult
	Tunnel       ProviderTunnelResult
	DispatchInfo RunDispatch
}

// SubmitProviderPool is the one-shot provider-pool dispatch surface. It performs
// a single queue admission, selects one candidate from the catalog, and dispatches
// the selected execution path — either tunnel/passthrough (OpenAI-compatible
// providers) or normalized RunEvent (Ollama/CLI/native). This method replaces
// the caller's need to choose between SubmitRun and SubmitProviderTunnel.
func (s *Service) SubmitProviderPool(ctx context.Context, req ProviderPoolDispatchRequest) (*ProviderPoolDispatchResult, error) {
	if s.queue == nil {
		return nil, fmt.Errorf("model queue is not configured")
	}

	candidates, returnedPolicy, err := s.resolveQueueCandidates(req.Run)
	if err != nil {
		return nil, err
	}

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

	long := req.Run.ContextClass == contextClassLong
	selected, queueReason, err := s.queue.admitWithReason(ctx, req.Run.ModelGroupKey, req.Run.Adapter, req.Run.Target, candidates, policy, s.resolveQueueCandidatesClosure(req.Run), long, req.Run.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 request hands it off to
	// the run/tunnel lifecycle.
	reservation := newQueueReservation(s.queue, selected)

	// Rewrite adapter and target for provider-pool dispatch: the winning candidate
	// carries the concrete adapter and served model name determined at selection time.
	adapter := req.Run.Adapter
	if selected.adapter != "" {
		adapter = selected.adapter
	}
	target := req.Run.Target
	if selected.servedTarget != "" {
		target = selected.servedTarget
	}

	switch selected.executionPath {
	case providerExecutionPathTunnel:
		return s.dispatchProviderPoolTunnel(req, adapter, target, selected, queueReason, reservation)

	case providerExecutionPathNormalized:
		runReq := req.Run
		if req.PrepareRun != nil {
			runReq, err = req.PrepareRun(runReq)
			if err != nil {
				reservation.release("prepare-run-error")
				return nil, err
			}
		}
		return s.dispatchProviderPoolRun(ctx, runReq, adapter, target, selected, queueReason, reservation)

	default:
		reservation.release("unknown-execution-path")
		return nil, fmt.Errorf("unknown execution path %q for provider-pool dispatch", selected.executionPath)
	}
}

// dispatchProviderPoolTunnel relays the selected candidate's raw provider
// request after provider-pool admission. The tunnel inherits the Run's
// identity, metadata, and long-context classification so passthrough dispatch
// keeps the same observability as the normalized path.
func (s *Service) dispatchProviderPoolTunnel(
	req ProviderPoolDispatchRequest,
	adapter, target string,
	selected *candidateNode,
	queueReason string,
	reservation *queueReservation,
) (*ProviderPoolDispatchResult, error) {
	tunnelReq := req.Tunnel
	tunnelReq.ProviderPool = true
	tunnelReq.ModelGroupKey = req.Run.ModelGroupKey
	tunnelReq.Adapter = adapter
	tunnelReq.Target = target
	tunnelReq.SessionID = req.Run.SessionID
	tunnelReq.MaxQueue = req.Run.MaxQueue
	tunnelReq.QueueTimeoutMS = req.Run.QueueTimeoutMS
	tunnelReq.Metadata = req.Run.Metadata
	tunnelReq.EstimatedInputTokens = req.Run.EstimatedInputTokens
	tunnelReq.ContextClass = req.Run.ContextClass

	// Apply pre-dispatch tunnel preparation (e.g. provider auth headers)
	// before buildProviderTunnelRequest so headers reach the wire request.
	if req.PrepareTunnel != nil {
		tunnelReqPrepared, prepErr := req.PrepareTunnel(tunnelReq)
		if prepErr != nil {
			reservation.release("prepare-tunnel-error")
			return nil, prepErr
		}
		tunnelReq = tunnelReqPrepared
	}

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

	reservation.track(runID)

	handle, err := s.openProviderTunnel(selected.entry, tunnelReqResolved, tunnelReq, queueReason, true, selected.providerID, selected.providerType, string(selected.executionPath))
	if err != nil {
		reservation.release("send-error")
		return nil, err
	}
	// The tunnel now owns the slot: it releases on the terminal END/ERROR
	// frame or on close.
	reservation.handOff()

	disp := handle.Dispatch()
	disp.ProviderID = selected.providerID
	disp.ProviderType = selected.providerType
	disp.ExecutionPath = string(selected.executionPath)
	disp.QueueReason = queueReason
	return &ProviderPoolDispatchResult{
		Path:         ProviderPoolPathTunnel,
		Tunnel:       handle,
		DispatchInfo: disp,
	}, nil
}

// dispatchProviderPoolRun is a helper that submits a normalized RunRequest after
// provider-pool admission. It tracks inflight, sends the RunRequest to the
// selected entry, and returns a RunResult wrapping the dispatch info.
func (s *Service) dispatchProviderPoolRun(
	ctx context.Context,
	req SubmitRunRequest,
	adapter, target string,
	selected *candidateNode,
	queueReason string,
	reservation *queueReservation,
) (*ProviderPoolDispatchResult, error) {
	req.Adapter = adapter
	req.Target = target

	runReq, runID, err := BuildRunRequest(req)
	if err != nil {
		reservation.release("build-error")
		return nil, err
	}

	// Track inflight before send so the event watcher can release the slot even
	// if a terminal event arrives before the Send call completes.
	reservation.track(runID)

	sub, err := s.subscribeRun(runID, selected.entry.NodeID, runReq.GetBackground())
	if err != nil {
		reservation.release("no-event-bus")
		return nil, err
	}

	if err := selected.entry.Client.Send(runReq); err != nil {
		reservation.release("send-error")
		sub.close()
		return nil, err
	}
	// The run now owns the slot: the event watcher releases it on the terminal
	// run event or on node disconnect.
	reservation.handOff()

	disp := RunDispatch{
		RunID:                runID,
		NodeID:               selected.entry.NodeID,
		NodeLabel:            nodeLabel(selected.entry),
		ModelGroupKey:        req.ModelGroupKey,
		Adapter:              runReq.GetAdapter(),
		Target:               runReq.GetTarget(),
		SessionID:            runReq.GetSessionId(),
		Background:           runReq.GetBackground(),
		TimeoutSec:           int(runReq.GetTimeoutSec()),
		EstimatedInputTokens: req.EstimatedInputTokens,
		ContextClass:         req.ContextClass,
		ProviderID:           selected.providerID,
		ProviderType:         selected.providerType,
		ExecutionPath:        string(selected.executionPath),
		QueueReason:          queueReason,
	}

	return &ProviderPoolDispatchResult{
		Path:         ProviderPoolPathNormalized,
		Run:          newRunHandle(disp, sub),
		DispatchInfo: disp,
	}, nil
}