iop/apps/edge/internal/openai/stream_gate_dispatcher.go

package openai

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/streamgate"
)

type openAIAdmissionKind string

const (
	openAIAdmissionRun    openAIAdmissionKind = "run"
	openAIAdmissionTunnel openAIAdmissionKind = "tunnel"
	openAIAdmissionPool   openAIAdmissionKind = "pool"
)

// openAIAttemptAdmission is built fresh for every recovery dispatch. The auth
// callback deliberately lives here rather than on ingress/rebuilt snapshots.
type openAIAttemptAdmission struct {
	kind      openAIAdmissionKind
	run       edgeservice.SubmitRunRequest
	tunnel    edgeservice.SubmitProviderTunnelRequest
	pool      edgeservice.ProviderPoolDispatchRequest
	authorize func(context.Context) (map[string]string, error)
}

func (a openAIAttemptAdmission) validate() error {
	switch a.kind {
	case openAIAdmissionRun:
		if a.run.ModelGroupKey == "" && a.run.Target == "" {
			return fmt.Errorf("OpenAI recovery run admission requires a model")
		}
	case openAIAdmissionTunnel:
		if a.tunnel.Path == "" {
			return fmt.Errorf("OpenAI recovery tunnel admission requires a path")
		}
	case openAIAdmissionPool:
		if !a.pool.Run.ProviderPool || a.pool.Run.ModelGroupKey == "" {
			return fmt.Errorf("OpenAI recovery pool admission is incomplete")
		}
	default:
		return fmt.Errorf("unknown OpenAI recovery admission kind")
	}
	return nil
}

type openAIAttemptAdmissionBuilder func(
	ctx context.Context,
	request streamgate.RebuiltRequest,
	body []byte,
) (openAIAttemptAdmission, error)

type openAIAttemptTransport struct {
	path         openAIAdmissionKind
	run          edgeservice.RunResult
	tunnel       edgeservice.ProviderTunnelResult
	usage        *openAIAttemptUsage
	usageBinding usageDispatchBinding
}

func (t *openAIAttemptTransport) bindUsage(dispatch edgeservice.RunDispatch) {
	if t == nil {
		return
	}
	t.usage = &openAIAttemptUsage{}
	t.usageBinding = newUsageDispatchBinding(dispatch, openAIAttemptResponseMode(t.path))
}

func openAIAttemptResponseMode(path openAIAdmissionKind) string {
	if path == openAIAdmissionTunnel {
		return responseModePassthrough
	}
	return responseModeNormalized
}

type openAIAttemptEventSourceFactory func(openAIAttemptTransport) (streamgate.NormalizedEventSource, error)

// openAIRecoveryAdmissionState carries only the sanitized classification of a
// recovery dispatch failure from the AttemptDispatcher to the release sink.
// Core intentionally hides raw host dispatcher errors, while the OpenAI host
// still has to preserve the public admission contract: an all-candidates
// capability rejection is a pre-dispatch 400 on initial, queued, and recovery
// admission alike.
type openAIRecoveryAdmissionState struct {
	mu                sync.Mutex
	candidateRejected bool
}

func (s *openAIRecoveryAdmissionState) record(err error) {
	if s == nil || !errors.Is(err, edgeservice.ErrProviderPoolCandidateRejected) {
		return
	}
	s.mu.Lock()
	s.candidateRejected = true
	s.mu.Unlock()
}

func (s *openAIRecoveryAdmissionState) rejected() bool {
	if s == nil {
		return false
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.candidateRejected
}

// openAIAttemptDispatcher adapts the three existing Edge admission surfaces
// to Core AttemptDispatcher. Provider/model/path values are never accepted
// from the rebuilder; they come exclusively from RunDispatch after admission.
type openAIAttemptDispatcher struct {
	service     runService
	store       *openAIRebuiltRequestStore
	build       openAIAttemptAdmissionBuilder
	eventSource openAIAttemptEventSourceFactory
	state       *openAIRecoveryAdmissionState
	usage       *openAIUsageRecorder
}

func newOpenAIAttemptDispatcher(
	service runService,
	store *openAIRebuiltRequestStore,
	build openAIAttemptAdmissionBuilder,
	eventSource openAIAttemptEventSourceFactory,
	usage ...*openAIUsageRecorder,
) (*openAIAttemptDispatcher, error) {
	if service == nil || store == nil || build == nil || eventSource == nil {
		return nil, fmt.Errorf("OpenAI attempt dispatcher dependencies are required")
	}
	dispatcher := &openAIAttemptDispatcher{
		service: service, store: store, build: build, eventSource: eventSource,
		state: &openAIRecoveryAdmissionState{},
	}
	if len(usage) > 0 {
		dispatcher.usage = usage[0]
	}
	return dispatcher, nil
}

func (d *openAIAttemptDispatcher) admissionState() *openAIRecoveryAdmissionState {
	return d.state
}

func (d *openAIAttemptDispatcher) DispatchAttempt(ctx context.Context, request streamgate.RebuiltRequest) (streamgate.AttemptBinding, error) {
	if err := ctx.Err(); err != nil {
		return streamgate.AttemptBinding{}, err
	}
	if err := request.Validate(); err != nil {
		return streamgate.AttemptBinding{}, err
	}
	lease, err := d.store.take(request.RequestRef())
	if err != nil {
		return streamgate.AttemptBinding{}, err
	}
	owned := true
	defer func() {
		if owned {
			lease.release()
		}
	}()

	body, err := lease.body()
	if err != nil {
		return streamgate.AttemptBinding{}, err
	}
	admission, err := d.build(ctx, request, body)
	if err != nil {
		return streamgate.AttemptBinding{}, err
	}
	if err := admission.validate(); err != nil {
		return streamgate.AttemptBinding{}, err
	}

	transport, dispatch, closeTransport, err := d.dispatch(ctx, admission)
	if err != nil {
		d.state.record(err)
		return streamgate.AttemptBinding{}, err
	}
	transport.bindUsage(dispatch)
	transportOwned := true
	defer func() {
		if transportOwned {
			closeTransport()
		}
	}()

	controller := &openAIAttemptController{
		service:        d.service,
		dispatch:       dispatch,
		closeTransport: closeTransport,
		lease:          lease,
		usageRecorder:  d.usage,
		usageBinding:   transport.usageBinding,
		usage:          transport.usage,
	}
	abortDispatched := func() {
		owned = false
		transportOwned = false
		_ = controller.AbortAttempt(context.Background())
	}
	source, err := d.eventSource(transport)
	if err != nil {
		abortDispatched()
		return streamgate.AttemptBinding{}, err
	}
	binding, err := streamgate.NewAttemptBinding(
		openAIStreamGateSafeToken("attempt", dispatch.RunID),
		actualOpenAIModel(dispatch),
		actualOpenAIProvider(dispatch),
		actualOpenAIExecutionPath(dispatch, transport.path),
		source,
		controller,
	)
	if err != nil {
		abortDispatched()
		return streamgate.AttemptBinding{}, err
	}
	owned = false
	transportOwned = false
	return binding, nil
}

func (d *openAIAttemptDispatcher) dispatch(
	ctx context.Context,
	admission openAIAttemptAdmission,
) (openAIAttemptTransport, edgeservice.RunDispatch, func(), error) {
	switch admission.kind {
	case openAIAdmissionRun:
		handle, err := d.service.SubmitRun(ctx, admission.run)
		if err != nil {
			return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, err
		}
		return openAIAttemptTransport{path: openAIAdmissionRun, run: handle}, handle.Dispatch(), handle.Close, nil

	case openAIAdmissionTunnel:
		request := admission.tunnel
		if admission.authorize != nil {
			headers, err := admission.authorize(ctx)
			if err != nil {
				return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, err
			}
			request.Headers = headers
		}
		handle, err := d.service.SubmitProviderTunnel(ctx, request)
		if err != nil {
			return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, err
		}
		return openAIAttemptTransport{path: openAIAdmissionTunnel, tunnel: handle}, handle.Dispatch(), handle.Close, nil

	case openAIAdmissionPool:
		request := admission.pool
		if admission.authorize != nil {
			previous := request.PrepareTunnel
			request.PrepareTunnel = func(tunnel edgeservice.SubmitProviderTunnelRequest) (edgeservice.SubmitProviderTunnelRequest, error) {
				headers, err := admission.authorize(ctx)
				if err != nil {
					return tunnel, err
				}
				tunnel.Headers = headers
				if previous != nil {
					return previous(tunnel)
				}
				return tunnel, nil
			}
		}
		result, err := d.service.SubmitProviderPool(ctx, request)
		if err != nil {
			return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, err
		}
		if result == nil {
			return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, fmt.Errorf("OpenAI provider pool returned no result")
		}
		switch result.Path {
		case edgeservice.ProviderPoolPathNormalized:
			if result.Run == nil {
				return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, fmt.Errorf("OpenAI provider pool normalized result is missing")
			}
			return openAIAttemptTransport{path: openAIAdmissionRun, run: result.Run}, result.DispatchInfo, result.Run.Close, nil
		case edgeservice.ProviderPoolPathTunnel:
			if result.Tunnel == nil {
				return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, fmt.Errorf("OpenAI provider pool tunnel result is missing")
			}
			return openAIAttemptTransport{path: openAIAdmissionTunnel, tunnel: result.Tunnel}, result.DispatchInfo, result.Tunnel.Close, nil
		default:
			return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, fmt.Errorf("OpenAI provider pool returned an unknown path")
		}
	default:
		return openAIAttemptTransport{}, edgeservice.RunDispatch{}, nil, fmt.Errorf("unknown OpenAI recovery admission kind")
	}
}

func actualOpenAIModel(dispatch edgeservice.RunDispatch) string {
	if model := strings.TrimSpace(dispatch.Target); model != "" {
		return model
	}
	return strings.TrimSpace(dispatch.ModelGroupKey)
}

func actualOpenAIProvider(dispatch edgeservice.RunDispatch) string {
	return strings.TrimSpace(dispatch.ProviderID)
}

func actualOpenAIExecutionPath(dispatch edgeservice.RunDispatch, path openAIAdmissionKind) string {
	if executionPath := strings.TrimSpace(dispatch.ExecutionPath); executionPath != "" {
		return executionPath
	}
	if path == openAIAdmissionTunnel {
		return string(edgeservice.ProviderPoolPathTunnel)
	}
	return string(edgeservice.ProviderPoolPathNormalized)
}

type openAIAttemptController struct {
	mu             sync.Mutex
	closed         bool
	service        runService
	dispatch       edgeservice.RunDispatch
	closeTransport func()
	lease          *openAIRebuiltLease
	usageRecorder  *openAIUsageRecorder
	usageBinding   usageDispatchBinding
	usage          *openAIAttemptUsage
}

func (c *openAIAttemptController) recordUsage() {
	if c == nil {
		return
	}
	c.usageRecorder.RecordAttempt(c.usageBinding, c.usage.snapshot())
}

// claimOwnership atomically transitions the controller to closed exactly once
// and hands the caller the transport closer and rebuilt lease so both
// AbortAttempt and CloseAttempt release ownership at most once.
func (c *openAIAttemptController) claimOwnership() (func(), *openAIRebuiltLease, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed {
		return nil, nil, false
	}
	c.closed = true
	closeTransport := c.closeTransport
	lease := c.lease
	c.closeTransport = nil
	c.lease = nil
	return closeTransport, lease, true
}

// AbortAttempt cancels the current provider run and releases its transport and
// rebuilt lease. It is idempotent and is the non-graceful teardown used for
// recovery replace, terminal error, and caller cancel.
func (c *openAIAttemptController) AbortAttempt(ctx context.Context) error {
	if c == nil {
		return nil
	}
	closeTransport, lease, claimed := c.claimOwnership()
	if !claimed {
		return nil
	}
	c.recordUsage()

	var cancelErr error
	if c.dispatch.RunID != "" {
		_, cancelErr = c.service.CancelRun(ctx, edgeservice.CancelRunRequest{
			NodeRef: c.dispatch.NodeID, RunID: c.dispatch.RunID,
		})
	}
	if closeTransport != nil {
		closeTransport()
	}
	lease.release()
	return cancelErr
}

// CloseAttempt releases the current provider transport and rebuilt lease
// without issuing a provider-side cancel. It is the graceful teardown used
// after a normal (successful) terminal, where the provider run already
// completed, and is idempotent with AbortAttempt through the shared claim.
func (c *openAIAttemptController) CloseAttempt(ctx context.Context) error {
	if c == nil {
		return nil
	}
	closeTransport, lease, claimed := c.claimOwnership()
	if !claimed {
		return nil
	}
	c.recordUsage()
	if closeTransport != nil {
		closeTransport()
	}
	lease.release()
	return nil
}

var _ streamgate.AttemptDispatcher = (*openAIAttemptDispatcher)(nil)
var _ streamgate.AttemptController = (*openAIAttemptController)(nil)