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

package service

import (
	"context"
	"fmt"
	"sync"
	"time"

	"google.golang.org/protobuf/types/known/structpb"

	edgenode "iop/apps/edge/internal/node"
	eventpkg "iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

type SubmitRunRequest struct {
	NodeRef        string
	RunID          string
	ModelGroupKey  string
	Adapter        string
	Target         string
	SessionID      string
	Workspace      string
	Prompt         string
	Input          map[string]any
	Background     bool
	TimeoutSec     int
	MaxQueue       int
	QueueTimeoutMS int
	Metadata       map[string]string
}

// RunDispatch describes a dispatched run in surface-neutral terms. It is the
// metadata any caller (console, HTTP, future RPC) needs after submission.
type RunDispatch struct {
	RunID         string
	NodeID        string
	NodeLabel     string
	ModelGroupKey string
	Adapter       string
	Target        string
	SessionID     string
	Background    bool
	TimeoutSec    int
}

// RunStream carries asynchronous events for a dispatched foreground run.
// Background runs leave both channels nil.
type RunStream struct {
	Events     <-chan *iop.RunEvent
	NodeEvents <-chan *iop.EdgeNodeEvent
}

// RunResult describes the interface for a submitted execution run, allowing
// surface-neutral consumption of its dispatch info and event stream.
type RunResult interface {
	Dispatch() RunDispatch
	Stream() RunStream
	Close()
	WaitTimeout() time.Duration
}

// RunHandle is the legacy combined surface kept for existing callers; it
// embeds the surface-neutral DTOs so HTTP/API code can consume RunDispatch
// or RunStream directly without depending on this struct.
type RunHandle struct {
	RunDispatch
	RunStream
	closeOnce sync.Once
	close     func()
}

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

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

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

func (h *RunHandle) Stream() RunStream {
	if h == nil {
		return RunStream{}
	}
	return h.RunStream
}

func (s *Service) SubmitRun(ctx context.Context, req SubmitRunRequest) (RunResult, error) {
	if req.ModelGroupKey != "" && s.queue != nil {
		return s.submitRunQueued(ctx, req)
	}
	return s.submitRunDirect(req)
}

func (s *Service) submitRunDirect(req SubmitRunRequest) (RunResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return nil, err
	}
	return s.dispatchToEntry(entry, req)
}

func (s *Service) submitRunQueued(ctx context.Context, req SubmitRunRequest) (RunResult, error) {
	candidates, policy, err := s.resolveQueueCandidates(req)
	if err != nil {
		return nil, err
	}

	entry, err := s.queue.admit(ctx, req.ModelGroupKey, req.Adapter, req.Target, candidates, policy)
	if err != nil {
		return nil, err
	}

	runReq, runID, err := BuildRunRequest(req)
	if err != nil {
		s.queue.releaseSlot(req.ModelGroupKey, entry.NodeID)
		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.
	s.queue.trackInflight(req.ModelGroupKey, runID, entry.NodeID)

	var runEvents <-chan *iop.RunEvent
	var unregisterRun func()
	var nodeEvents <-chan *iop.EdgeNodeEvent
	var unregisterNode func()
	if !runReq.GetBackground() {
		if s.events == nil {
			s.queue.releaseRun(runID, "no-event-bus")
			return nil, fmt.Errorf("event bus is not configured")
		}
		runEvents, unregisterRun = s.events.SubscribeRun(runID, 4096)
		nodeEvents, unregisterNode = s.events.SubscribeNode(entry.NodeID, 16)
	}

	if err := entry.Client.Send(runReq); err != nil {
		s.queue.releaseRun(runID, "send-error")
		if unregisterRun != nil {
			unregisterRun()
		}
		if unregisterNode != nil {
			unregisterNode()
		}
		return nil, err
	}

	return &RunHandle{
		RunDispatch: RunDispatch{
			RunID:         runID,
			NodeID:        entry.NodeID,
			NodeLabel:     nodeLabel(entry),
			ModelGroupKey: req.ModelGroupKey,
			Adapter:       runReq.GetAdapter(),
			Target:        runReq.GetTarget(),
			SessionID:     runReq.GetSessionId(),
			Background:    runReq.GetBackground(),
			TimeoutSec:    int(runReq.GetTimeoutSec()),
		},
		RunStream: RunStream{
			Events:     runEvents,
			NodeEvents: nodeEvents,
		},
		close: func() {
			if unregisterRun != nil {
				unregisterRun()
			}
			if unregisterNode != nil {
				unregisterNode()
			}
		},
	}, nil
}

// resolveQueueCandidates returns candidate nodes filtered by adapter/target capability,
// each paired with per-node capacity, plus the group policy derived from the
// request route policy or adapter config.
func (s *Service) resolveQueueCandidates(req SubmitRunRequest) ([]candidateNode, groupPolicy, error) {
	if req.NodeRef != "" {
		entry, err := s.ResolveNode(req.NodeRef)
		if err != nil {
			return nil, groupPolicy{}, err
		}
		cap := defaultNodeCapacity
		if s.nodeStore != nil {
			if rec, ok := s.nodeStore.FindByID(entry.NodeID); ok {
				res := resolveAdapterForNode(rec, req.Adapter, req.Target)
				if !res.supported {
					msg := fmt.Sprintf("node %q does not support adapter %q target %q", entry.NodeID, req.Adapter, req.Target)
					if res.ambiguous {
						msg = fmt.Sprintf("node %q: adapter %q is ambiguous (multiple enabled instances); use an instance key", entry.NodeID, req.Adapter)
					}
					return nil, groupPolicy{}, fmt.Errorf("%s", msg)
				}
				cap = res.capacity
			}
		}
		policy := groupPolicyFromRequestOrStore(req, s.nodeStore, []*edgenode.NodeEntry{entry})
		return []candidateNode{{entry: entry, capacity: cap}}, policy, nil
	}

	all := s.registry.All()
	if len(all) == 0 {
		return nil, groupPolicy{}, fmt.Errorf("no nodes connected")
	}

	var candidates []candidateNode
	for _, entry := range all {
		cap := defaultNodeCapacity
		if s.nodeStore != nil {
			rec, ok := s.nodeStore.FindByID(entry.NodeID)
			if ok {
				res := resolveAdapterForNode(rec, req.Adapter, req.Target)
				if !res.supported {
					continue
				}
				cap = res.capacity
			}
		}
		candidates = append(candidates, candidateNode{entry: entry, capacity: cap})
	}
	if len(candidates) == 0 {
		return nil, groupPolicy{}, fmt.Errorf("no nodes support adapter %q target %q", req.Adapter, req.Target)
	}

	entries := make([]*edgenode.NodeEntry, len(candidates))
	for i, c := range candidates {
		entries[i] = c.entry
	}
	policy := groupPolicyFromRequestOrStore(req, s.nodeStore, entries)
	return candidates, policy, nil
}

// adapterResolution holds the resolved capacity and queue policy for a single
// node/adapter combination. ambiguous is true when a type-name lookup (e.g.,
// "ollama") matches 2+ enabled instances on that node, mirroring the Node
// router's exact-instance-key/ambiguity contract.
type adapterResolution struct {
	supported      bool
	ambiguous      bool
	capacity       int
	maxQueue       int
	queueTimeoutMS int
}

func positiveOr(v, fallback int) int {
	if v > 0 {
		return v
	}
	return fallback
}

// resolveAdapterForNode determines whether a node can handle adapterType/target
// and computes the per-node capacity and queue policy fields.
//
// Resolution order:
//  1. Exact instance Name match across OllamaInstances / VllmInstances /
//     OpenAICompatInstances (instance-key route, matching Node router priority).
//  2. Type-name route (e.g. "ollama"): supported only when exactly 1 enabled
//     instance of that type exists. 2+ enabled instances → ambiguous (fail,
//     same semantics as Node router ambiguity error). 0 instances → fail-open
//     for legacy/unconfigured nodes.
//  3. "cli": capability gated by CLI.Enabled and profile name in target.
//  4. Default (unknown adapter type): fail-open.
func resolveAdapterForNode(rec *edgenode.NodeRecord, adapterType, target string) adapterResolution {
	if rec == nil {
		return adapterResolution{supported: true, capacity: defaultNodeCapacity}
	}
	concurrencyFallback := positiveOr(rec.Runtime.Concurrency, defaultNodeCapacity)

	// Exact instance Name match (highest priority).
	for _, inst := range rec.Adapters.OllamaInstances {
		if inst.Name == adapterType {
			return adapterResolution{
				supported:      inst.Enabled,
				capacity:       positiveOr(inst.Capacity, concurrencyFallback),
				maxQueue:       inst.MaxQueue,
				queueTimeoutMS: inst.QueueTimeoutMS,
			}
		}
	}
	for _, inst := range rec.Adapters.VllmInstances {
		if inst.Name == adapterType {
			return adapterResolution{
				supported:      inst.Enabled,
				capacity:       positiveOr(inst.Capacity, concurrencyFallback),
				maxQueue:       inst.MaxQueue,
				queueTimeoutMS: inst.QueueTimeoutMS,
			}
		}
	}
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if inst.Name == adapterType {
			return adapterResolution{
				supported:      inst.Enabled,
				capacity:       positiveOr(inst.Capacity, concurrencyFallback),
				maxQueue:       inst.MaxQueue,
				queueTimeoutMS: inst.QueueTimeoutMS,
			}
		}
	}

	// Type-name route.
	switch adapterType {
	case "ollama":
		return resolveTypeRoute(rec, ollamaEnabledInstances(rec), concurrencyFallback)
	case "vllm":
		return resolveTypeRoute(rec, vllmEnabledInstances(rec), concurrencyFallback)
	case "openai_compat":
		return resolveTypeRoute(rec, openAICompatEnabledInstances(rec), concurrencyFallback)
	case "cli":
		// CLI has no named multi-instance model; capability is gated by profile.
		if !rec.Adapters.CLI.Enabled && len(rec.Adapters.CLI.Profiles) == 0 {
			// No CLI config at all → fail-open for legacy/unconfigured nodes.
			return adapterResolution{supported: true, capacity: concurrencyFallback}
		}
		if !rec.Adapters.CLI.Enabled {
			return adapterResolution{supported: false}
		}
		if target == "" {
			return adapterResolution{supported: len(rec.Adapters.CLI.Profiles) > 0, capacity: concurrencyFallback}
		}
		_, ok := rec.Adapters.CLI.Profiles[target]
		return adapterResolution{supported: ok, capacity: concurrencyFallback}
	default:
		return adapterResolution{supported: true, capacity: concurrencyFallback}
	}
}

// instanceFields holds the capacity/policy fields extracted from a single adapter instance.
type instanceFields struct {
	capacity, maxQueue, queueTimeoutMS int
}

func ollamaEnabledInstances(rec *edgenode.NodeRecord) []instanceFields {
	var out []instanceFields
	for _, inst := range rec.Adapters.OllamaInstances {
		if inst.Enabled {
			out = append(out, instanceFields{inst.Capacity, inst.MaxQueue, inst.QueueTimeoutMS})
		}
	}
	return out
}

func vllmEnabledInstances(rec *edgenode.NodeRecord) []instanceFields {
	var out []instanceFields
	for _, inst := range rec.Adapters.VllmInstances {
		if inst.Enabled {
			out = append(out, instanceFields{inst.Capacity, inst.MaxQueue, inst.QueueTimeoutMS})
		}
	}
	return out
}

func openAICompatEnabledInstances(rec *edgenode.NodeRecord) []instanceFields {
	var out []instanceFields
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if inst.Enabled {
			out = append(out, instanceFields{inst.Capacity, inst.MaxQueue, inst.QueueTimeoutMS})
		}
	}
	return out
}

// resolveTypeRoute applies Node-router-compatible type-name resolution:
//   - 0 enabled instances → fail-open (legacy / unconfigured)
//   - 1 enabled instance → use its capacity and policy
//   - 2+ enabled instances → ambiguous (reject)
func resolveTypeRoute(rec *edgenode.NodeRecord, enabled []instanceFields, concurrencyFallback int) adapterResolution {
	switch len(enabled) {
	case 0:
		return adapterResolution{supported: true, capacity: concurrencyFallback}
	case 1:
		f := enabled[0]
		return adapterResolution{
			supported:      true,
			capacity:       positiveOr(f.capacity, concurrencyFallback),
			maxQueue:       f.maxQueue,
			queueTimeoutMS: f.queueTimeoutMS,
		}
	default:
		return adapterResolution{supported: false, ambiguous: true}
	}
}

func groupPolicyFromRequestOrStore(req SubmitRunRequest, store *edgenode.NodeStore, entries []*edgenode.NodeEntry) groupPolicy {
	if req.MaxQueue > 0 || req.QueueTimeoutMS > 0 {
		maxQueue := req.MaxQueue
		if maxQueue <= 0 {
			maxQueue = defaultGroupMaxQueue
		}
		queueTimeout := defaultQueueTimeout
		if req.QueueTimeoutMS > 0 {
			queueTimeout = time.Duration(req.QueueTimeoutMS) * time.Millisecond
		}
		return groupPolicy{
			maxQueue:     maxQueue,
			queueTimeout: queueTimeout,
		}
	}
	return groupPolicyFromStore(store, entries, req.Adapter, req.Target)
}

// groupPolicyFromStore derives queue policy from the first resolved candidate node.
func groupPolicyFromStore(store *edgenode.NodeStore, entries []*edgenode.NodeEntry, adapterType, target string) groupPolicy {
	if store != nil {
		for _, e := range entries {
			rec, ok := store.FindByID(e.NodeID)
			if !ok {
				continue
			}
			res := resolveAdapterForNode(rec, adapterType, target)
			if !res.supported {
				continue
			}
			if res.maxQueue > 0 || res.queueTimeoutMS > 0 {
				p := groupPolicy{
					maxQueue:     positiveOr(res.maxQueue, defaultGroupMaxQueue),
					queueTimeout: time.Duration(positiveOr(res.queueTimeoutMS, 0)) * time.Millisecond,
				}
				if p.queueTimeout <= 0 {
					p.queueTimeout = defaultQueueTimeout
				}
				return p
			}
		}
	}
	return groupPolicy{maxQueue: defaultGroupMaxQueue, queueTimeout: defaultQueueTimeout}
}

func (s *Service) dispatchToEntry(entry *edgenode.NodeEntry, req SubmitRunRequest) (RunResult, error) {
	runReq, runID, err := BuildRunRequest(req)
	if err != nil {
		return nil, err
	}

	var runEvents <-chan *iop.RunEvent
	var unregisterRun func()
	var nodeEvents <-chan *iop.EdgeNodeEvent
	var unregisterNode func()
	if !runReq.GetBackground() {
		if s.events == nil {
			return nil, fmt.Errorf("event bus is not configured")
		}
		runEvents, unregisterRun = s.events.SubscribeRun(runID, 4096)
		nodeEvents, unregisterNode = s.events.SubscribeNode(entry.NodeID, 16)
	}

	if err := entry.Client.Send(runReq); err != nil {
		if unregisterRun != nil {
			unregisterRun()
		}
		if unregisterNode != nil {
			unregisterNode()
		}
		return nil, err
	}

	return &RunHandle{
		RunDispatch: RunDispatch{
			RunID:         runID,
			NodeID:        entry.NodeID,
			NodeLabel:     nodeLabel(entry),
			ModelGroupKey: req.ModelGroupKey,
			Adapter:       runReq.GetAdapter(),
			Target:        runReq.GetTarget(),
			SessionID:     runReq.GetSessionId(),
			Background:    runReq.GetBackground(),
			TimeoutSec:    int(runReq.GetTimeoutSec()),
		},
		RunStream: RunStream{
			Events:     runEvents,
			NodeEvents: nodeEvents,
		},
		close: func() {
			if unregisterRun != nil {
				unregisterRun()
			}
			if unregisterNode != nil {
				unregisterNode()
			}
		},
	}, nil
}

type CancelRunRequest struct {
	NodeRef   string
	RunID     string
	Adapter   string
	Target    string
	SessionID string
}

func BuildCancelRunRequest(req CancelRunRequest) *iop.CancelRequest {
	return &iop.CancelRequest{
		RunId:     req.RunID,
		Adapter:   req.Adapter,
		Target:    req.Target,
		SessionId: NormalizeSessionID(req.SessionID),
		Action:    iop.CancelAction_CANCEL_ACTION_CANCEL_RUN,
	}
}

func (s *Service) CancelRun(_ context.Context, req CancelRunRequest) (CommandResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return CommandResult{}, err
	}
	cancelReq := BuildCancelRunRequest(req)
	if err := entry.Client.Send(cancelReq); err != nil {
		return CommandResult{}, err
	}
	return CommandResult{
		NodeID:    entry.NodeID,
		NodeLabel: nodeLabel(entry),
		SessionID: cancelReq.GetSessionId(),
	}, nil
}

type TerminateSessionRequest struct {
	NodeRef   string
	Adapter   string
	Target    string
	SessionID string
}

// CommandResult is the surface-neutral acknowledgement for one-shot node
// commands (terminate-session, future control RPCs).
type CommandResult struct {
	NodeID    string
	NodeLabel string
	SessionID string
}

// TerminateSessionResult is retained as an alias for backward compatibility.
type TerminateSessionResult = CommandResult

func (s *Service) TerminateSession(_ context.Context, req TerminateSessionRequest) (TerminateSessionResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return TerminateSessionResult{}, err
	}

	sessionID := NormalizeSessionID(req.SessionID)
	cancelReq := &iop.CancelRequest{
		Adapter:   req.Adapter,
		Target:    req.Target,
		SessionId: sessionID,
		Action:    iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
	}
	if err := entry.Client.Send(cancelReq); err != nil {
		return TerminateSessionResult{}, err
	}
	return TerminateSessionResult{
		NodeID:    entry.NodeID,
		NodeLabel: nodeLabel(entry),
		SessionID: sessionID,
	}, nil
}

func NewRunID() string {
	return fmt.Sprintf("manual-%d", time.Now().UnixNano())
}

func IsNodeDisconnected(event *iop.EdgeNodeEvent) bool {
	return event.GetType() == eventpkg.TypeNodeDisconnected
}

func BuildRunRequest(req SubmitRunRequest) (*iop.RunRequest, string, error) {
	inputMap := make(map[string]any, len(req.Input)+1)
	for k, v := range req.Input {
		inputMap[k] = v
	}
	if req.Prompt != "" {
		if _, ok := inputMap["prompt"]; !ok {
			inputMap["prompt"] = req.Prompt
		}
	}
	input, err := structpb.NewStruct(inputMap)
	if err != nil {
		return nil, "", err
	}
	runID := req.RunID
	if runID == "" {
		runID = NewRunID()
	}
	metadata := make(map[string]string, len(req.Metadata))
	for k, v := range req.Metadata {
		metadata[k] = v
	}
	return &iop.RunRequest{
		RunId:       runID,
		Adapter:     req.Adapter,
		Target:      req.Target,
		Workspace:   req.Workspace,
		SessionId:   NormalizeSessionID(req.SessionID),
		SessionMode: iop.RunSessionMode_RUN_SESSION_MODE_CREATE_IF_MISSING,
		Background:  req.Background,
		Input:       input,
		TimeoutSec:  int32(normalizeTimeoutSec(req.TimeoutSec)),
		Metadata:    metadata,
	}, runID, nil
}