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

package service

import (
	"fmt"
	"strconv"

	runtime "iop/packages/go/execution"
	iop "iop/proto/gen/iop"
)

const recoveryHandoffConfirmed = "confirmed"

type receivedTerminalDisposition uint8

const (
	receivedTerminalUntracked receivedTerminalDisposition = iota
	receivedTerminalAccepted
	receivedTerminalRejected
)

func isTerminalRunEvent(e *iop.RunEvent) bool {
	t := e.GetType()
	return t == "complete" || t == "error" || t == "cancelled"
}

// releaseLease frees the resources held by one lease, exactly once, whatever the
// cause: a send failure, a normalized terminal run event, a tunnel END/ERROR
// frame, a tunnel close, or a node disconnect. All of them converge here, and the
// lease map — mutated only under m.mu — decides which caller performs the single
// transition into released. A lease id that is already gone is a no-op, so a
// losing racer can never decrement a counter that now belongs to another request.
func (m *modelQueueManager) releaseLease(leaseID uint64, reason string) {
	if leaseID == 0 {
		return
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.leases[leaseID]; !ok {
		return
	}
	if !m.releaseLeaseLocked(leaseID) {
		return
	}
	// Pump every group, not just the one that held the lease: the freed slot
	// belongs to a provider resource that other model groups queue against, and
	// the earliest waiter for it may be in any of them.
	m.pumpAllLocked()
}

// releaseLeaseLocked performs the reserved/tracked → released transition and
// frees the slot the lease holds, without pumping the queue. Returns true when
// this call is the one that performed the transition. Must be called with m.mu
// held.
func (m *modelQueueManager) releaseLeaseLocked(leaseID uint64) bool {
	lease, ok := m.leases[leaseID]
	if !ok {
		return false
	}
	delete(m.leases, leaseID)
	if lease.runID != "" {
		delete(m.leaseByRun, lease.runID)
	}
	lease.state = leaseStateReleased
	m.decrementSlotLocked(lease.groupKey, lease.nodeID, lease.providerID, lease.long)
	return true
}

// receivedHealthEvidence is the fully validated Node health evidence carried by
// one typed response-stalled terminal. It contains no caller-controlled fields.
type receivedHealthEvidence struct {
	providerHealth string
	sequence       uint64
}

func parseReceivedHealthEvidence(runID string, failure *iop.ExecutionFailure) (receivedHealthEvidence, bool) {
	if failure == nil || failure.GetCode() != string(runtime.FailureCodeResponseStalled) || !failure.GetRetryable() {
		return receivedHealthEvidence{}, false
	}
	metadata := failure.GetMetadata()
	if metadata["failure_code"] != string(runtime.FailureCodeResponseStalled) ||
		metadata["attempt_fence"] != "confirmed" ||
		metadata["run_id"] != runID || metadata["attempt_id"] != runID ||
		metadata["adapter"] == "" || metadata["target"] == "" {
		return receivedHealthEvidence{}, false
	}
	sequence, err := strconv.ParseUint(metadata["health_observation_seq"], 10, 64)
	if err != nil || sequence == 0 {
		return receivedHealthEvidence{}, false
	}
	health := metadata["provider_health"]
	classification := metadata["liveness_classification"]
	switch {
	case health == string(runtime.ProviderStatusUnavailable) && classification == string(runtime.ProviderUnhealthy):
	case health == string(runtime.ProviderStatusAvailable) && classification == string(runtime.RequestStalled):
	case health == string(runtime.ProviderStatusUnknown) && classification == string(runtime.HealthUnknown):
	default:
		return receivedHealthEvidence{}, false
	}
	return receivedHealthEvidence{providerHealth: health, sequence: sequence}, true
}

func annotateRecoveryHandoff(failure *iop.ExecutionFailure, envelopeMetadata *map[string]string, providerID, providerHealth string) {
	if failure.Metadata == nil {
		failure.Metadata = make(map[string]string)
	}
	failure.Metadata["provider_id"] = providerID
	failure.Metadata["provider_health"] = providerHealth
	failure.Metadata["recovery_handoff"] = recoveryHandoffConfirmed
	if envelopeMetadata == nil {
		return
	}
	if *envelopeMetadata == nil {
		*envelopeMetadata = make(map[string]string)
	}
	(*envelopeMetadata)["provider_id"] = providerID
	(*envelopeMetadata)["provider_health"] = providerHealth
	(*envelopeMetadata)["recovery_handoff"] = recoveryHandoffConfirmed
}

// applyReceivedHealthEvidenceLocked sequence-fences one fully bound terminal.
// Every accepted observation advances the high-water mark. Only unavailable
// lowers effective provider health; available/unknown stall observations never
// recover an already unavailable provider.
func (m *modelQueueManager) applyReceivedHealthEvidenceLocked(lease *providerLease, evidence receivedHealthEvidence) providerHealthObservation {
	observation := providerHealthObservation{
		source: "stall", evidenceHealth: evidence.providerHealth, decision: "inconclusive",
	}
	if lease == nil || lease.providerID == "" || lease.adapter == "" || lease.target == "" {
		observation.decision = "rejected_binding"
		return observation
	}
	key := providerRuntimeHealthKey{
		nodeID: lease.nodeID, generation: lease.generation, providerID: lease.providerID,
	}
	overlay := m.runtimeHealth[key]
	observation.fromHealth = runtimeOverlayHealth(overlay)
	observation.toHealth = observation.fromHealth
	if overlay != nil && evidence.sequence <= overlay.observationSeq {
		observation.decision = "rejected_stale"
		return observation
	}
	if overlay == nil {
		overlay = &providerRuntimeHealthOverlay{}
		m.runtimeHealth[key] = overlay
	}
	overlay.observationSeq = evidence.sequence
	if evidence.providerHealth == string(runtime.ProviderStatusUnavailable) {
		overlay.adapter = lease.adapter
		overlay.target = lease.target
		overlay.unavailable = true
	}
	observation.decision = "applied"
	observation.toHealth = runtimeOverlayHealth(overlay)
	observation.stateChanged = observation.fromHealth != observation.toHealth
	return observation
}

// settleReceivedTerminal validates authoritative reception identity against the
// immutable lease before any correctness state changes. A current terminal
// releases once even when its optional health evidence is missing or rejected.
// A mismatched node/generation is rejected and cannot release another owner's
// lease. For accepted bound stall evidence, handoff annotation, any fresh overlay
// transition, release, and queue pumping all occur under m.mu.
func (m *modelQueueManager) settleReceivedTerminal(nodeID string, generation uint64, runID string, failure *iop.ExecutionFailure, envelopeMetadata *map[string]string) receivedTerminalDisposition {
	if runID == "" {
		return receivedTerminalUntracked
	}
	m.mu.Lock()

	leaseID, tracked := m.leaseByRun[runID]
	if !tracked {
		m.mu.Unlock()
		return receivedTerminalUntracked
	}
	lease := m.leases[leaseID]
	if lease == nil {
		delete(m.leaseByRun, runID)
		m.mu.Unlock()
		return receivedTerminalUntracked
	}
	if nodeID == "" || generation == 0 || lease.nodeID != nodeID || lease.generation != generation {
		m.mu.Unlock()
		return receivedTerminalRejected
	}

	var observation *providerHealthObservation
	if evidence, ok := parseReceivedHealthEvidence(runID, failure); ok {
		metadata := failure.GetMetadata()
		if lease.providerID != "" && metadata["adapter"] == lease.adapter && metadata["target"] == lease.target {
			// Handoff confirms authoritative reception, immutable lease binding,
			// and the local attempt fence. Sequence freshness governs only the
			// provider-wide overlay; an out-of-order terminal still carries its
			// request-local handoff and still releases its own lease.
			annotateRecoveryHandoff(failure, envelopeMetadata, lease.providerID, evidence.providerHealth)
			result := m.applyReceivedHealthEvidenceLocked(lease, evidence)
			observation = &result
		} else {
			result := providerHealthObservation{source: "stall", evidenceHealth: evidence.providerHealth, decision: "rejected_binding"}
			observation = &result
		}
	}

	if m.releaseLeaseLocked(leaseID) {
		m.pumpAllLocked()
	}
	m.mu.Unlock()
	m.observeProviderHealth(observation)
	return receivedTerminalAccepted
}

// resolveCurrentProbeProviderLocked resolves CAPABILITIES evidence against the
// authoritative current Node provider catalog. Runtime overlays are a health
// projection, not an identity source: a healthy sibling provider with the same
// adapter/target must make the evidence ambiguous as well. Must be called with
// m.mu held.
func (m *modelQueueManager) resolveCurrentProbeProviderLocked(nodeID, adapter, target string) (string, bool) {
	if m.store == nil {
		return "", false
	}
	record, ok := m.store.FindByID(nodeID)
	if !ok || record == nil {
		return "", false
	}

	matchedProviderID := ""
	for _, provider := range record.Providers {
		if provider.ID == "" || providerAdapterKey(provider) != adapter || !providerCanServe(provider, target) {
			continue
		}
		if matchedProviderID != "" {
			return "", false
		}
		matchedProviderID = provider.ID
	}
	return matchedProviderID, matchedProviderID != ""
}

// applyProviderProbeEvidence offers one CAPABILITIES health observation to the
// current provider catalog. The exact adapter/target must identify one and only
// one current-generation provider. A strictly newer available observation is
// recorded even if the provider is already effectively available, so a delayed
// lower-sequence terminal cannot later mark it unavailable. Only an actual
// unavailable-to-available transition pumps the queue and reports recovery.
func (m *modelQueueManager) applyProviderProbeEvidence(nodeID string, generation uint64, adapter, target string, status runtime.ProviderStatus, sequence uint64, isCurrentOwner func() bool) bool {
	observation := providerHealthObservation{source: "probe", evidenceHealth: string(status), decision: "inconclusive"}
	if nodeID == "" || generation == 0 || adapter == "" || target == "" ||
		status != runtime.ProviderStatusAvailable || sequence == 0 {
		m.observeProviderHealth(&observation)
		return false
	}
	m.mu.Lock()
	if isCurrentOwner != nil && !isCurrentOwner() {
		m.mu.Unlock()
		observation.decision = "rejected_binding"
		m.observeProviderHealth(&observation)
		return false
	}

	providerID, ok := m.resolveCurrentProbeProviderLocked(nodeID, adapter, target)
	if !ok {
		m.mu.Unlock()
		observation.decision = "rejected_ambiguous"
		m.observeProviderHealth(&observation)
		return false
	}
	key := providerRuntimeHealthKey{nodeID: nodeID, generation: generation, providerID: providerID}
	overlay := m.runtimeHealth[key]
	observation.fromHealth = runtimeOverlayHealth(overlay)
	observation.toHealth = observation.fromHealth
	if overlay != nil && sequence <= overlay.observationSeq {
		m.mu.Unlock()
		observation.decision = "rejected_stale"
		m.observeProviderHealth(&observation)
		return false
	}
	if overlay == nil {
		overlay = &providerRuntimeHealthOverlay{}
		m.runtimeHealth[key] = overlay
	}
	recovered := overlay.unavailable && overlay.adapter == adapter && overlay.target == target
	overlay.observationSeq = sequence
	if overlay.unavailable && !recovered {
		m.mu.Unlock()
		observation.decision = "rejected_binding"
		observation.toHealth = runtimeOverlayHealth(overlay)
		m.observeProviderHealth(&observation)
		return false
	}
	overlay.adapter = adapter
	overlay.target = target
	overlay.unavailable = false
	if recovered {
		m.pumpAllLocked()
	}
	observation.decision = "applied"
	observation.toHealth = runtimeOverlayHealth(overlay)
	observation.stateChanged = observation.fromHealth != observation.toHealth
	m.mu.Unlock()
	m.observeProviderHealth(&observation)
	return recovered
}

func runtimeOverlayHealth(overlay *providerRuntimeHealthOverlay) string {
	if overlay != nil && overlay.unavailable {
		return string(runtime.ProviderStatusUnavailable)
	}
	return string(runtime.ProviderStatusAvailable)
}

// fenceNodeGenerationLocked fences the disconnected connection identified by
// (nodeID, generation): it settles leases through the exactly-once release path
// so each provider resource counter is returned per lease, and marks matching
// provider resources orphan. Under a generation-scoped fence a strictly newer
// reconnect's leases and resources are left untouched; under a whole-node fence
// (generation 0) it also zeroes any residual counters seeded by fixtures
// without a backing lease. Returns whether any lease was settled so the caller
// can tell an authoritative fence from a stale no-op. Must be called with m.mu
// held.
func (m *modelQueueManager) fenceNodeGenerationLocked(nodeID string, generation uint64) bool {
	fenceAll := generation == 0
	settledLease := m.settleLeasesForNodeLocked(nodeID, generation)

	for key, res := range m.resources {
		if key.nodeID != nodeID {
			continue
		}
		if !fenceAll && res.generation > generation {
			continue
		}
		if fenceAll {
			res.inFlight = 0
			res.longInFlight = 0
		}
		res.orphan = true
	}
	for key := range m.runtimeHealth {
		if key.nodeID == nodeID && (fenceAll || key.generation <= generation) {
			delete(m.runtimeHealth, key)
		}
	}
	return settledLease
}

// resolveAndPumpAllLocked rebuilds every queued item's candidate universe
// through its live resolver (or filters orphaned and fenced-node snapshot
// candidates when no resolver is attached) and then runs the global pump.
// Terminal no-candidate items are settled immediately so the pump only deals
// with dispatchable or temporarily-blocked waiters. excludeNodeID, excludeGen
// and fenceAll describe the disconnect context so snapshot items that lack a
// live resolver are also purged of the fenced node — the same signal the old
// candidate-snapshot removal loop used — while provider-pool items with a
// resolver get a fresh resolution against current state. This is the
// authoritative disconnect settlement pass: surviving-provider fallback,
// terminal unavailable, and normal dispatch all resolve in the same
// deterministic global pass, independent of event bus delivery. Must be called
// with m.mu held.
func (m *modelQueueManager) resolveAndPumpAllLocked(excludeNodeID string, excludeGen uint64, fenceAll bool) {
	for _, group := range m.groups {
		for _, item := range group.queue {
			if item.resolveCandidates != nil {
				candidates, outcome, resolveErr := m.resolveQueuedCandidatesLocked(item)
				if outcome == resolveNoCandidates {
					// Terminal no-candidate: remove immediately, before the pump,
					// so the dispatch pass only sees items that can actually run.
					m.removeQueuedItemLocked(group, item)
					select {
					case item.waitCh <- admitResult{err: fmt.Errorf("model group %q: %w", group.key, errProviderUnavailable)}:
					default:
					}
					continue
				}
				if outcome == resolveResolverError {
					// Leave the item queued for a later pump; the resolver fault
					// is recoverable and must not block the dispatch pass.
					continue
				}
				if outcome == resolveTerminalError {
					m.removeQueuedItemLocked(group, item)
					select {
					case item.waitCh <- admitResult{err: resolveErr}:
					default:
					}
					continue
				}
				if candidates != nil {
					item.candidates = candidates
				}
				continue
			}
			// Snapshot-candidate path: filter orphaned resources AND the fenced
			// node itself so a pump dispatch cannot hand them to an offline
			// connection. A strictly newer reconnect on the same node is
			// preserved (generation-scoped fence).
			filtered := make([]candidateNode, 0, len(item.candidates))
			for _, c := range item.candidates {
				if c.entry == nil {
					filtered = append(filtered, c)
					continue
				}
				if excludeNodeID != "" && (fenceAll || c.generation <= excludeGen) && c.entry.NodeID == excludeNodeID {
					continue
				}
				key := providerResourceKey{nodeID: c.entry.NodeID, providerID: c.providerID}
				if r, ok := m.resources[key]; ok && r.orphan {
					continue
				}
				filtered = append(filtered, c)
			}
			item.candidates = filtered
		}
	}
	// One global pump after every group's candidates are refreshed, so a waiter
	// can only be handed a provider that is still live, and a waiter in any
	// group can pick up the remaining ones.
	m.pumpAllLocked()
}

// releaseNode fences the disconnected connection identified by (nodeID,
// generation): it drops that connection's leases, zeroes the matching provider
// resources to offline, rebuilds every queued item's candidate universe through
// its live resolver, and runs one global pump so surviving-provider fallback,
// terminal unavailable, and normal dispatch all resolve in the same pass. A
// generation of 0 fences the node unconditionally (untracked/legacy disconnect).
// A specific generation fences only resources whose current owner generation is
// that one or older: a stale callback for an already-superseded connection
// cannot zero the slot a live reconnect now holds, so its live resource and
// events are left unchanged.
func (m *modelQueueManager) releaseNode(nodeID string, generation uint64, reason string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	fenceAll := generation == 0

	// Fence the generation: settle leases and mark matching resources orphan.
	settledLease := m.fenceNodeGenerationLocked(nodeID, generation)

	// A generation-scoped callback that fenced neither a lease nor a resource is
	// a stale close for an already-superseded connection: the live owner's state
	// and queued candidate lists are authoritative, so leave them untouched.
	if !fenceAll && !settledLease {
		return
	}

	// A whole-node fence (generation 0) clears any residual legacy/provider slot
	// counters this node still holds — legacy (nodeID-only) fixtures and any
	// counter a lease did not account for. Provider-pool counters live on the
	// resource state already marked orphan above.
	if fenceAll {
		for _, group := range m.groups {
			for slot := range group.inflight {
				if colonIdx := findLastColon(slot); colonIdx > 0 {
					if slot[:colonIdx] == nodeID {
						delete(group.inflight, slot)
						delete(group.longInflight, slot)
					}
				} else if slot == nodeID {
					delete(group.inflight, slot)
					delete(group.longInflight, slot)
				}
			}
		}
	}

	// Rebuild every queued item's candidate universe through the live resolver
	// and run one global pump so that surviving-provider fallback, terminal
	// unavailable, and normal dispatch all resolve in the same deterministic
	// pass. This replaces the previous candidate-snapshot removal loop: items
	// with a live resolver get re-evaluated against the current store/catalog/
	// registry state (including orphan filtering), and items without a live
	// resolver have the fenced node and orphaned resources purged from their
	// candidate lists, so the pump sees only live targets and can settle
	// terminal no-candidate items directly.
	m.resolveAndPumpAllLocked(nodeID, generation, fenceAll)
}

// activateNode is the reconnect counterpart to releaseNode: it restores the
// provider resources owned by an accepted connection identified by (nodeID,
// generation) to an available state for that generation, then rebuilds every
// queued item's candidate universe through its live resolver and runs one global
// pump. A disconnect marked those resources orphan and stranded the waiters
// queued against them; the reconnect alone re-activates the resources and
// re-dispatches the waiters in global enqueue order, with no new request, config
// refresh, or lease release to trigger the pump.
//
// No disconnect context is passed to the rebuild pass (excludeNodeID empty,
// fenceAll false), so nothing is purged: the reconnected node's candidates are
// made live and dispatched, and every other group's waiters are re-evaluated
// against current state in the same deterministic pass.
//
// isCurrentOwner is the ownership linearization gate. It is evaluated under m.mu,
// in the same critical section that mutates the resources and pumps, and it is
// serialized against releaseNode (which also holds m.mu). A ready callback whose
// connection a disconnect settled — unregistered its entry and released its
// leases — between the transport handshake and this activation therefore observes
// the disconnect and is a no-op: the check and the activation cannot straddle the
// disconnect the way a check performed before m.mu was acquired could. The
// callback itself reads the registry (m.mu → registry-lock order, matching the
// pump's live resolver); the registry lock is never held across m.mu, so no lock
// inversion is introduced. A nil callback is treated as always-current for
// untracked/legacy callers and fixtures.
func (m *modelQueueManager) activateNode(nodeID string, generation uint64, isCurrentOwner func() bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if isCurrentOwner != nil && !isCurrentOwner() {
		// A stale/superseded generation: the disconnect that removed this
		// connection already settled its resources under this same lock, and
		// re-activating them here would resurrect a slot the live owner (or none)
		// now holds.
		return
	}
	m.activateNodeGenerationLocked(nodeID, generation)
	m.resolveAndPumpAllLocked("", 0, false)
}

// settleLeasesForNodeLocked returns the resources held by the leases a disconnect
// fences on a node, routing each through the exactly-once releaseLeaseLocked
// decrement path so per-lease provider counters are returned — never zeroed
// wholesale. A generation of 0 fences every lease on the node (untracked/legacy
// disconnect); a specific generation fences that owner and any older one, leaving
// a strictly newer reconnect's leases — and the shared resource counters they
// still hold — untouched. It reports whether it settled any lease so the caller
// can tell an authoritative fence from a stale no-op. Must be called with m.mu
// held.
func (m *modelQueueManager) settleLeasesForNodeLocked(nodeID string, generation uint64) bool {
	fenceAll := generation == 0
	// Collect first: releaseLeaseLocked mutates m.leases, so it cannot run inside
	// the range over the same map.
	matching := make([]uint64, 0, len(m.leases))
	for id, lease := range m.leases {
		if lease.nodeID != nodeID {
			continue
		}
		if !fenceAll && lease.generation > generation {
			continue
		}
		matching = append(matching, id)
	}
	for _, id := range matching {
		m.releaseLeaseLocked(id)
	}
	return len(matching) > 0
}

// releaseRun releases the lease owning a terminated run and dispatches the next
// queued item if one is waiting. Safe to call for an unknown or already-released
// run: the lease index resolves nothing and the call is a no-op.
func (m *modelQueueManager) releaseRun(runID, reason string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	leaseID, ok := m.leaseByRun[runID]
	if !ok {
		return
	}
	if m.leases[leaseID] == nil {
		delete(m.leaseByRun, runID)
		return
	}
	if !m.releaseLeaseLocked(leaseID) {
		return
	}
	m.pumpAllLocked()
}

// releaseSlot decrements the in-flight count and tries to dispatch the next
// queued item. Used when admit-path I/O fails after the slot was reserved.
// For provider-pool dispatches (nodeID:providerID slot), pass providerID.
// For legacy dispatches, pass empty providerID so nodeID is used directly.
func (m *modelQueueManager) releaseSlot(groupKey, nodeID string, providerID ...string) {
	var pid string
	if len(providerID) > 0 {
		pid = providerID[0]
	}
	m.releaseSlotWithLong(groupKey, nodeID, pid, false)
}

// releaseSlotWithLong releases a reserved slot, additionally freeing a long-context
// slot when longSlot is true. Used by admit-path failure/race handling where the
// caller knows whether a long slot was reserved.
func (m *modelQueueManager) releaseSlotWithLong(groupKey, nodeID, providerID string, longSlot bool) {
	m.mu.Lock()
	m.releaseSlotLocked(groupKey, nodeID, providerID, longSlot)
	m.mu.Unlock()
}

func (m *modelQueueManager) releaseSlotLocked(groupKey, nodeID string, providerID string, longSlot bool) {
	if group := m.decrementSlotLocked(groupKey, nodeID, providerID, longSlot); group == nil {
		return
	}
	m.pumpAllLocked()
}

// decrementSlotLocked frees the resources a reservation holds without pumping the
// queue, returning the group so the caller can decide whether to dispatch. Must
// be called with m.mu held.
func (m *modelQueueManager) decrementSlotLocked(groupKey, nodeID, providerID string, longSlot bool) *modelQueueGroup {
	group, ok := m.groups[groupKey]
	if !ok {
		return nil
	}
	// Use provider-aware slot key for provider-pool dispatches.
	slot := nodeID
	if providerID != "" {
		slot = nodeID + ":" + providerID
		key := providerResourceKey{nodeID: nodeID, providerID: providerID}
		if res, ok := m.resources[key]; ok {
			res.release(longSlot)
			if res.orphan && res.inFlight == 0 && res.longInFlight == 0 {
				delete(m.resources, key)
			}
		} else {
			if group.inflight[slot] > 0 {
				group.inflight[slot]--
			}
			if longSlot && group.longInflight[slot] > 0 {
				group.longInflight[slot]--
			}
		}
	} else {
		if group.inflight[slot] > 0 {
			group.inflight[slot]--
		}
		if longSlot && group.longInflight[slot] > 0 {
			group.longInflight[slot]--
		}
	}
	return group
}