iop/apps/node/internal/node/run_handler.go

package node

import (
	"context"
	"errors"
	"fmt"
	"time"

	"go.uber.org/zap"

	"iop/apps/node/internal/store"
	"iop/apps/node/internal/transport"
	runtime "iop/packages/go/execution"
	iop "iop/proto/gen/iop"
)

// OnRunRequest handles an incoming RunRequest from a transport Session.
func (n *Node) OnRunRequest(ctx context.Context, sess *transport.Session, req *iop.RunRequest) error {
	n.logger.Info("run request received", zap.String("run_id", req.GetRunId()), zap.String("adapter", req.GetAdapter()), zap.String("target", req.GetTarget()))

	rr := runRequestFromProto(req)
	printEdgeMessage(n.out, rr.Input)

	if err := n.validateRunStallTimeout(sess, req, &rr); err != nil {
		return err
	}

	n.configSetMu.RLock()
	configLocked := true
	defer func() {
		if configLocked {
			n.configSetMu.RUnlock()
		}
	}()

	spec, adapter, err := n.router.ResolveAdapter(ctx, rr)
	if err != nil {
		n.sendPreExecuteError(sess, req.GetRunId(), req.GetSessionId(), req.GetBackground(), n.nodeID, err.Error())
		return fmt.Errorf("node: resolve: %w", err)
	}

	var caps runtime.Capabilities
	if c, capsErr := adapter.Capabilities(ctx); capsErr == nil {
		caps = c
	}
	admission := n.admissionFor(spec.Adapter, caps)

	ticket, err := admission.acquire()
	if err != nil {
		n.logger.Warn("run admission rejected",
			zap.String("run_id", spec.RunID),
			zap.String("adapter", spec.Adapter),
			zap.String("reason", string(admissionRejectReason(err))),
		)
		n.rejectRun(ctx, sess, spec, err)
		return fmt.Errorf("node: run %s: %w", spec.RunID, err)
	}

	if err := n.store.InsertRun(ctx, store.RunRecord{
		RunID:      spec.RunID,
		Adapter:    spec.Adapter,
		Target:     spec.Target,
		SessionID:  normalizeSessionID(spec.SessionID),
		Background: spec.Background,
		Status:     "running",
		CreatedAt:  time.Now(),
	}); err != nil {
		n.logger.Warn("store: insert run", zap.String("run_id", spec.RunID), zap.Error(err))
	}

	// Session listeners supply their connection-lifetime context. Direct callers
	// retain the context they provided.
	execCtx, cancel := context.WithCancel(ctx)
	if spec.TimeoutSec > 0 {
		execCtx, cancel = context.WithTimeout(ctx, time.Duration(spec.TimeoutSec)*time.Second)
	}

	h := &runHandle{
		runID:     spec.RunID,
		adapter:   spec.Adapter,
		target:    spec.Target,
		sessionID: normalizeSessionID(spec.SessionID),
		cancel:    cancel,
		done:      make(chan struct{}),
	}
	n.runs.register(h)

	configLocked = false
	n.configSetMu.RUnlock()

	var sender protoSender = noopSender{}
	var seq healthObservationSequencer
	if sess != nil && sess.IsAlive() {
		sender = sess
		seq = sess
	}
	probe := healthProbeFor(adapter, caps.AdapterName, caps.InstanceKey, spec.Target)

	run := func() error {
		return n.executeNormalizedAttempt(ctx, execCtx, cancel, adapter, spec, ticket, h, sender, probe, seq)
	}

	if spec.Background {
		go func() { _ = run() }()
		return nil
	}
	return run()
}

// admissionFor returns an admissionManager with only the per-adapter gate for
// the given adapter instance key. Node-wide global admission gating has been
// removed; each adapter controls its own concurrency via Capabilities().MaxConcurrency.
func (n *Node) admissionFor(adapterKey string, caps runtime.Capabilities) *admissionManager {
	return &admissionManager{
		adapter: n.adapterGateFor(adapterKey, caps),
	}
}

// adapterGateFor returns the shared gate for the given adapter instance key,
// creating one from caps if not yet seen. MaxConcurrency <= 0 means the adapter gate is unlimited.
func (n *Node) adapterGateFor(adapterKey string, caps runtime.Capabilities) *fifoGate {
	n.adapterGatesMu.Lock()
	defer n.adapterGatesMu.Unlock()
	g, ok := n.adapterGates[adapterKey]
	if !ok {
		g = newFifoGate(caps.MaxConcurrency)
		n.adapterGates[adapterKey] = g
	}
	return g
}

// rejectRun records a rejected run in the store and sends an error RunEvent to
// the session so Edge can observe the rejection through the normal event stream.
// The rejection message carries the concrete reason (concurrency_unavailable) extracted from the admission error.
func (n *Node) rejectRun(ctx context.Context, sess *transport.Session, spec runtime.ExecutionSpec, admitErr error) {
	errMsg := admissionMessage(admitErr)

	// Since we reject immediately without queueing, insert as a terminal rejected record.
	if err := n.store.InsertRun(ctx, store.RunRecord{
		RunID:      spec.RunID,
		Adapter:    spec.Adapter,
		Target:     spec.Target,
		SessionID:  normalizeSessionID(spec.SessionID),
		Background: spec.Background,
		Status:     "rejected",
		CreatedAt:  time.Now(),
	}); err != nil {
		n.logger.Warn("store: insert rejected run", zap.String("run_id", spec.RunID), zap.Error(err))
	}
	if err := n.store.CompleteRun(ctx, spec.RunID, "rejected", errMsg); err != nil {
		n.logger.Warn("store: complete rejected run", zap.String("run_id", spec.RunID), zap.Error(err))
	}

	if sess != nil && sess.IsAlive() {
		re := &iop.RunEvent{
			RunId:      spec.RunID,
			Type:       string(runtime.EventTypeError),
			Error:      errMsg,
			Timestamp:  time.Now().UnixNano(),
			SessionId:  normalizeSessionID(spec.SessionID),
			Background: spec.Background,
			NodeId:     n.nodeID,
		}
		if err := sess.Send(re); err != nil {
			n.logger.Warn("session: send reject event", zap.String("run_id", spec.RunID), zap.Error(err))
		}
	}
}

// sendPreExecuteError sends an error RunEvent when an error occurs before
// adapter execution (e.g. ResolveAdapter failure). No store record is needed
// because the run never reached execution. This ensures Edge can observe the
// failure and avoid inflight slot leaks.
func (n *Node) sendPreExecuteError(sess *transport.Session, runID, sessionID string, background bool, nodeID, errMsg string) {
	if sess != nil && sess.IsAlive() {
		re := &iop.RunEvent{
			RunId:      runID,
			Type:       string(runtime.EventTypeError),
			Error:      errMsg,
			Timestamp:  time.Now().UnixNano(),
			SessionId:  normalizeSessionID(sessionID),
			Background: background,
			NodeId:     nodeID,
		}
		if err := sess.Send(re); err != nil {
			n.logger.Warn("session: send pre-execute error event", zap.String("run_id", runID), zap.Error(err))
		}
	}
}

func (n *Node) completeRun(spec runtime.ExecutionSpec, execErr error) {
	status := "completed"
	errMsg := ""
	if execErr != nil {
		if errors.Is(execErr, runtime.ErrRunCancelled) {
			status = "cancelled"
		} else {
			status = "failed"
			errMsg = execErr.Error()
		}
		n.logger.Warn("run ended", zap.String("run_id", spec.RunID), zap.String("status", status), zap.Error(execErr))
	}
	if err := n.store.CompleteRun(context.Background(), spec.RunID, status, errMsg); err != nil {
		n.logger.Warn("store: complete run", zap.String("run_id", spec.RunID), zap.Error(err))
	}
}

// synthAndEmitTerminal queues a terminal event when the adapter returned
// without emitting one. The caller flushes it after local admission release, so
// Edge can observe run completion without over-dispatching back into Node.
func (n *Node) synthAndEmitTerminal(ctx context.Context, sink *terminalDeferringSink, spec runtime.ExecutionSpec, execErr error) error {
	event := runtime.RuntimeEvent{
		RunID:     spec.RunID,
		Timestamp: time.Now(),
	}
	switch {
	case errors.Is(execErr, runtime.ErrRunCancelled):
		event.Type = runtime.EventTypeCancelled
	case execErr != nil:
		event.Type = runtime.EventTypeError
		event.Error = execErr.Error()
		event.Failure = runtime.FailureFromError(execErr)
	default:
		event.Type = runtime.EventTypeComplete
		event.Message = "adapter completed without terminal event"
	}
	if err := sink.Emit(ctx, event); err != nil {
		return fmt.Errorf("synthesize terminal event: %w", err)
	}
	return nil
}