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

// Package node is the core IOP Node service. It implements
// transport.Handler and orchestrates routing → adapter execution.
package node

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"sync"

	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/structpb"

	"iop/apps/node/internal/adapters"
	"iop/apps/node/internal/router"
	"iop/apps/node/internal/runtime"
	"iop/apps/node/internal/store"
	"iop/apps/node/internal/transport"
	iop "iop/proto/gen/iop"
)

// Node implements transport.Handler and coordinates the full execution pipeline.
type Node struct {
	nodeID           string
	router           runtime.Router
	store            *store.Store
	runs             *runManager
	globalGate       *fifoGate // node-wide concurrency safety guard across all adapters
	adapterGatesMu   sync.Mutex
	adapterGates     map[string]*fifoGate // per adapter-key concurrency safety guard
	out              io.Writer
	logger           *zap.Logger
	currentConfigSet *adapters.ConfigSet
	configSetMu      sync.RWMutex
}

// New creates a Node. It satisfies transport.Handler.
// globalConcurrency is the node-wide maximum for simultaneous executions across
// all adapters (>= 1 enforces limit; <= 0 means unlimited at the node level).
// Per-adapter MaxConcurrency from Capabilities is applied independently on top.
// out receives console debug output; pass os.Stdout for production, io.Discard in tests.
func New(
	nodeID string,
	router runtime.Router,
	st *store.Store,
	globalConcurrency int,
	out io.Writer,
	logger *zap.Logger,
	initialConfigSet *adapters.ConfigSet,
) *Node {
	if out == nil {
		out = os.Stdout
	}
	return &Node{
		nodeID:           nodeID,
		router:           router,
		store:            st,
		runs:             newRunManager(),
		globalGate:       newFifoGate(globalConcurrency),
		adapterGates:     make(map[string]*fifoGate),
		out:              out,
		logger:           logger,
		currentConfigSet: initialConfigSet,
	}
}

// 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 := runtime.RunRequest{
		RunID:       req.GetRunId(),
		Adapter:     req.GetAdapter(),
		Target:      req.GetTarget(),
		SessionID:   req.GetSessionId(),
		SessionMode: sessionModeFromProto(req.GetSessionMode()),
		Background:  req.GetBackground(),
		Workspace:   req.GetWorkspace(),
		Policy:      structAsMap(req.GetPolicy()),
		Input:       structAsMap(req.GetInput()),
		TimeoutSec:  int(req.GetTimeoutSec()),
		Metadata:    req.GetMetadata(),
	}
	printEdgeMessage(n.out, rr.Input)

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

	spec, adapter, err := n.router.ResolveAdapter(ctx, rr)
	if err != nil {
		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)

	// Acquire safety capacity ticket. Since we no longer maintain a Node-local FIFO queue,
	// if concurrency is full, we reject immediately.
	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)
	}

	// Record the request as running since it is admitted immediately without queueing.
	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))
	}

	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{}
	if sess != nil && sess.IsAlive() {
		sender = sess
	}

	sink := &sessionSink{
		sess:       sender,
		out:        n.out,
		nodeID:     n.nodeID,
		sessionID:  normalizeSessionID(spec.SessionID),
		background: spec.Background,
	}
	runSink := &terminalDeferringSink{inner: sink}

	run := func() error {
		released := false
		releaseTicket := func() {
			if !released {
				ticket.release()
				released = true
			}
		}
		defer releaseTicket()
		defer cancel()
		defer n.runs.deregister(spec.RunID)
		defer close(h.done)

		execErr := adapter.Execute(execCtx, spec, runSink)
		releaseTicket()
		n.completeRun(spec, execErr)
		if flushErr := runSink.Flush(context.Background()); flushErr != nil {
			n.logger.Warn("session: flush terminal events", zap.String("run_id", spec.RunID), zap.Error(flushErr))
			if execErr == nil {
				return flushErr
			}
		}
		return execErr
	}

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

// admissionFor returns an admissionManager pairing the node-wide global gate
// with the per-adapter gate for the given adapter instance key.
func (n *Node) admissionFor(adapterKey string, caps runtime.Capabilities) *admissionManager {
	return &admissionManager{
		global:  n.globalGate,
		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))
		}
	}
}

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))
	}
}

// OnCancel cancels a running execution or terminates an adapter session.
func (n *Node) OnCancel(_ context.Context, _ *transport.Session, req *iop.CancelRequest) error {
	n.logger.Info("cancel request", zap.String("run_id", req.GetRunId()), zap.String("action", req.GetAction().String()))

	switch cancelActionFromProto(req.GetAction()) {
	case runtime.CancelActionTerminateSession:
		adapter, err := n.router.LookupAdapter(req.GetAdapter())
		if err != nil {
			return fmt.Errorf("node: %w", err)
		}
		terminator, ok := adapter.(runtime.SessionTerminator)
		if !ok {
			return fmt.Errorf("node: adapter %q does not support session termination", req.GetAdapter())
		}
		return terminator.TerminateSession(context.Background(), req.GetTarget(), normalizeSessionID(req.GetSessionId()))
	default:
		n.runs.cancelRun(req.GetRunId())
		return nil
	}
}

// OnConfigRefresh handles a config refresh request from edge.
func (n *Node) OnConfigRefresh(ctx context.Context, _ *transport.Session, req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error) {
	n.logger.Info("config refresh request received", zap.String("request_id", req.GetRequestId()))

	if req.GetConfig() == nil {
		if len(req.GetChangedPaths()) == 0 {
			return &iop.NodeConfigRefreshResponse{
				RequestId: req.GetRequestId(),
				Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED,
			}, nil
		}
		return &iop.NodeConfigRefreshResponse{
			RequestId:            req.GetRequestId(),
			Status:               iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_RESTART_REQUIRED,
			RestartRequiredPaths: req.GetChangedPaths(),
		}, nil
	}

	nextSet, err := adapters.BuildConfigSet(req.GetConfig(), n.logger)
	if err != nil {
		n.logger.Error("config refresh: failed to build next config set", zap.Error(err))
		return &iop.NodeConfigRefreshResponse{
			RequestId: req.GetRequestId(),
			Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED,
			Error:     fmt.Sprintf("build config set: %s", err.Error()),
		}, nil
	}

	n.configSetMu.Lock()
	defer n.configSetMu.Unlock()

	// Runtime tuning. Concurrency is applied live but deferred until after
	// adapter registry start/swap succeeds so that a failed refresh preserves
	// the previous admission state.
	var nextConcurrency int
	var hasRuntimeConcurrencyChange bool
	if rc := req.GetConfig().GetRuntime(); rc != nil {
		// Pre-compute the concurrency change without mutating the gate yet.
		nextConcurrency = int(rc.GetConcurrency())
		hasRuntimeConcurrencyChange = nextConcurrency != n.globalGate.currentCapacity()
	}

	diff := adapters.DiffConfigSets(n.currentConfigSet, nextSet)

	if len(diff.Added) == 0 && len(diff.Updated) == 0 && len(diff.Removed) == 0 {
		n.logger.Info("config refresh: no adapter changes detected")
		if n.currentConfigSet != nil {
			n.currentConfigSet.Runtime = nextSet.Runtime
		}
		// Apply runtime concurrency since there are no adapter changes to fail.
		if hasRuntimeConcurrencyChange {
			n.globalGate.updateCapacity(nextConcurrency)
		}
		return &iop.NodeConfigRefreshResponse{
			RequestId: req.GetRequestId(),
			Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED,
		}, nil
	}

	n.logger.Info("config refresh: adapter diff calculated",
		zap.Strings("added", diff.Added),
		zap.Strings("updated", diff.Updated),
		zap.Strings("removed", diff.Removed),
	)

	// Start the new registry
	if err := nextSet.Registry.Start(ctx); err != nil {
		n.logger.Error("config refresh: failed to start new adapters", zap.Error(err))
		return &iop.NodeConfigRefreshResponse{
			RequestId: req.GetRequestId(),
			Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED,
			Error:     fmt.Sprintf("start next registry: %s", err.Error()),
		}, nil
	}

	// Live swap registry in router
	if mr, ok := n.router.(router.MutableRouter); ok {
		mr.SetRegistry(nextSet.Registry)
	} else {
		n.logger.Warn("router does not implement MutableRouter, cannot swap registry")
	}

	// Update existing adapter gates capacity
	n.adapterGatesMu.Lock()
	for key, gate := range n.adapterGates {
		if adapter, ok := nextSet.Registry.Get(key); ok {
			if caps, err := adapter.Capabilities(ctx); err == nil {
				gate.updateCapacity(caps.MaxConcurrency)
			}
		}
	}
	n.adapterGatesMu.Unlock()

	// Only apply runtime concurrency after all adapter operations succeed.
	if hasRuntimeConcurrencyChange {
		n.logger.Info("config refresh: applying runtime concurrency after successful adapter swap",
			zap.Int("previous", n.globalGate.currentCapacity()),
			zap.Int("next", nextConcurrency),
		)
		n.globalGate.updateCapacity(nextConcurrency)
	}

	// Registry.Stop stops every lifecycle adapter in the old registry. Pre-swap
	// runs still hold old adapter instances, so the old registry stop is deferred
	// until those runs drain; in-flight executions then finish on their original
	// adapter snapshot while new requests already route to nextSet.
	oldSet := n.currentConfigSet
	pendingRuns := n.runs.snapshotActive()

	if len(pendingRuns) == 0 {
		if oldSet != nil && oldSet.Registry != nil {
			if err := oldSet.Registry.Stop(context.Background()); err != nil {
				n.logger.Warn("config refresh: failed to stop old adapters", zap.Error(err))
			}
		}
	} else if oldSet != nil && oldSet.Registry != nil {
		n.logger.Info("config refresh: deferring old registry stop until active runs drain",
			zap.Int("pending_runs", len(pendingRuns)),
		)
		go func() {
			waitHandles(context.Background(), pendingRuns)
			if err := oldSet.Registry.Stop(context.Background()); err != nil {
				n.logger.Warn("config refresh: failed to stop old adapters after drain", zap.Error(err))
			}
		}()
	}

	n.currentConfigSet = nextSet

	return &iop.NodeConfigRefreshResponse{
		RequestId: req.GetRequestId(),
		Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED,
	}, nil
}

// OnCommandRequest handles an incoming node command from a transport Session.
func (n *Node) OnCommandRequest(ctx context.Context, sess *transport.Session, req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
	n.logger.Info("command request",
		zap.String("request_id", req.GetRequestId()),
		zap.String("type", req.GetType().String()),
		zap.String("adapter", req.GetAdapter()),
		zap.String("target", req.GetTarget()),
	)

	cmdType, ok := protoCommandTypeToDomain(req.GetType())
	if !ok {
		return n.commandErrorResponse(req, fmt.Sprintf("node: unsupported command type %q", req.GetType().String())), nil
	}

	execCtx := ctx
	if req.GetTimeoutSec() > 0 {
		var cancel context.CancelFunc
		execCtx, cancel = context.WithTimeout(ctx, time.Duration(req.GetTimeoutSec())*time.Second)
		defer cancel()
	}

	switch cmdType {
	case runtime.CommandTypeCapabilities:
		return n.handleCapabilitiesCommand(execCtx, req), nil
	case runtime.CommandTypeTransportStatus:
		return n.handleTransportStatusCommand(sess, req), nil
	default:
		return n.dispatchAdapterCommand(execCtx, req), nil
	}
}

func (n *Node) handleCapabilitiesCommand(ctx context.Context, req *iop.NodeCommandRequest) *iop.NodeCommandResponse {
	adapter, err := n.router.LookupAdapter(req.GetAdapter())
	if err != nil {
		return n.commandErrorResponse(req, fmt.Sprintf("node: %s", err.Error()))
	}
	caps, err := adapter.Capabilities(ctx)
	if err != nil {
		return n.commandErrorResponse(req, err.Error())
	}

	targets := append([]string(nil), caps.Targets...)
	providerStatus := caps.ProviderStatus
	providerDetail := ""

	if prober, ok := adapter.(runtime.ProviderProber); ok {
		probeRes, err := prober.ProbeProvider(ctx, req.GetTarget())
		if err != nil {
			providerStatus = runtime.ProviderStatusUnavailable
			providerDetail = err.Error()
		} else {
			providerStatus = probeRes.Status
			providerDetail = probeRes.Detail
			if len(probeRes.Targets) > 0 {
				targets = append([]string(nil), probeRes.Targets...)
			}
		}
	}

	sort.Strings(targets)

	n.adapterGatesMu.Lock()
	gate, ok := n.adapterGates[req.GetAdapter()]
	n.adapterGatesMu.Unlock()
	inFlight := 0
	queued := 0
	if ok {
		inFlight = gate.activeCount()
	}

	result := map[string]string{
		"adapter":         caps.AdapterName,
		"instance_key":    caps.InstanceKey,
		"targets":         strings.Join(targets, ","),
		"max_concurrency": strconv.Itoa(caps.MaxConcurrency),
		"provider_status": string(runtime.NormalizeProviderStatus(providerStatus)),
		"capacity":        strconv.Itoa(caps.MaxConcurrency),
		"in_flight":       strconv.Itoa(inFlight),
		"queued":          strconv.Itoa(queued),
	}
	if providerDetail != "" {
		result["provider_detail"] = providerDetail
	}

	providerSnapshot := &iop.ProviderSnapshot{
		Adapter:  req.GetAdapter(),
		Status:   string(runtime.NormalizeProviderStatus(providerStatus)),
		Capacity: int32(caps.MaxConcurrency),
		InFlight: int32(inFlight),
		Queued:   int32(queued),
	}

	return &iop.NodeCommandResponse{
		RequestId:         req.GetRequestId(),
		Type:              req.GetType(),
		Adapter:           req.GetAdapter(),
		Target:            req.GetTarget(),
		SessionId:         req.GetSessionId(),
		Result:            result,
		ProviderSnapshots: []*iop.ProviderSnapshot{providerSnapshot},
	}
}

func sessionConnected(sess *transport.Session) (alive bool) {
	if sess == nil {
		return false
	}
	defer func() {
		if r := recover(); r != nil {
			alive = false
		}
	}()
	return sess.IsAlive()
}

func (n *Node) handleTransportStatusCommand(sess *transport.Session, req *iop.NodeCommandRequest) *iop.NodeCommandResponse {
	connected := "false"
	state := "disconnected"
	if sessionConnected(sess) {
		connected = "true"
		state = "connected"
	}
	result := map[string]string{
		"node_id":    n.nodeID,
		"connected":  connected,
		"state":      state,
		"adapter":    req.GetAdapter(),
		"target":     req.GetTarget(),
		"session_id": normalizeSessionID(req.GetSessionId()),
	}
	return &iop.NodeCommandResponse{
		RequestId: req.GetRequestId(),
		Type:      req.GetType(),
		Adapter:   req.GetAdapter(),
		Target:    req.GetTarget(),
		SessionId: normalizeSessionID(req.GetSessionId()),
		Result:    result,
	}
}

func (n *Node) dispatchAdapterCommand(ctx context.Context, req *iop.NodeCommandRequest) *iop.NodeCommandResponse {
	adapter, err := n.router.LookupAdapter(req.GetAdapter())
	if err != nil {
		return n.commandErrorResponse(req, fmt.Sprintf("node: %s", err.Error()))
	}
	handler, ok := adapter.(runtime.CommandHandler)
	if !ok {
		return n.commandErrorResponse(req, fmt.Sprintf("node: adapter %q does not support commands", req.GetAdapter()))
	}
	domainReq := toDomainCommandRequest(req)
	domainResp, err := handler.HandleCommand(ctx, domainReq)
	if err != nil {
		return n.commandErrorResponse(req, err.Error())
	}
	return toProtoCommandResponse(domainResp)
}

func (n *Node) commandErrorResponse(req *iop.NodeCommandRequest, msg string) *iop.NodeCommandResponse {
	return &iop.NodeCommandResponse{
		RequestId: req.GetRequestId(),
		Type:      req.GetType(),
		Adapter:   req.GetAdapter(),
		Target:    req.GetTarget(),
		SessionId: req.GetSessionId(),
		Error:     msg,
	}
}

// protoCommandTypeToDomain maps a proto NodeCommandType to its runtime
// CommandType. The second return is false for UNSPECIFIED or unknown values
// so callers can reject unsupported commands before dispatching.
func protoCommandTypeToDomain(t iop.NodeCommandType) (runtime.CommandType, bool) {
	switch t {
	case iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS:
		return runtime.CommandTypeUsageStatus, true
	case iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES:
		return runtime.CommandTypeCapabilities, true
	case iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST:
		return runtime.CommandTypeSessionList, true
	case iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS:
		return runtime.CommandTypeTransportStatus, true
	case iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API:
		return runtime.CommandTypeOllamaAPI, true
	default:
		return "", false
	}
}

func domainCommandTypeToProto(t runtime.CommandType) iop.NodeCommandType {
	switch t {
	case runtime.CommandTypeUsageStatus:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS
	case runtime.CommandTypeCapabilities:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES
	case runtime.CommandTypeSessionList:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST
	case runtime.CommandTypeTransportStatus:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS
	case runtime.CommandTypeOllamaAPI:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API
	default:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_UNSPECIFIED
	}
}

func toDomainCommandRequest(req *iop.NodeCommandRequest) runtime.CommandRequest {
	cmdType, _ := protoCommandTypeToDomain(req.GetType())
	return runtime.CommandRequest{
		RequestID:  req.GetRequestId(),
		Type:       cmdType,
		Adapter:    req.GetAdapter(),
		Target:     req.GetTarget(),
		SessionID:  normalizeSessionID(req.GetSessionId()),
		TimeoutSec: int(req.GetTimeoutSec()),
		Metadata:   req.GetMetadata(),
	}
}

func toProtoCommandResponse(resp runtime.CommandResponse) *iop.NodeCommandResponse {
	out := &iop.NodeCommandResponse{
		RequestId: resp.RequestID,
		Type:      domainCommandTypeToProto(resp.Type),
		Adapter:   resp.Adapter,
		Target:    resp.Target,
		SessionId: resp.SessionID,
		Result:    resp.Result,
	}
	if resp.UsageStatus != nil {
		out.UsageStatus = &iop.AgentUsageStatus{
			RawOutput:       resp.UsageStatus.RawOutput,
			DailyLimit:      resp.UsageStatus.DailyLimit,
			DailyResetTime:  resp.UsageStatus.DailyResetTime,
			WeeklyLimit:     resp.UsageStatus.WeeklyLimit,
			WeeklyResetTime: resp.UsageStatus.WeeklyResetTime,
			Metadata:        resp.UsageStatus.Metadata,
		}
	}
	return out
}

type protoSender interface {
	Send(m proto.Message) error
}

type noopSender struct{}

func (noopSender) Send(proto.Message) error { return nil }

// terminalDeferringSink holds terminal events until Node-local admission has
// released its slot. Edge uses terminal run events to advance queued work, so
// emitting them before the local slot is free can over-dispatch back into Node.
type terminalDeferringSink struct {
	inner runtime.EventSink

	mu        sync.Mutex
	deferring bool
	deferred  []runtime.RuntimeEvent
}

func (s *terminalDeferringSink) Emit(ctx context.Context, event runtime.RuntimeEvent) error {
	s.mu.Lock()
	if s.deferring || isTerminalRuntimeEvent(event.Type) {
		s.deferring = true
		s.deferred = append(s.deferred, event)
		s.mu.Unlock()
		return nil
	}
	s.mu.Unlock()
	return s.inner.Emit(ctx, event)
}

func (s *terminalDeferringSink) Flush(ctx context.Context) error {
	s.mu.Lock()
	events := append([]runtime.RuntimeEvent(nil), s.deferred...)
	s.deferred = nil
	s.deferring = false
	s.mu.Unlock()

	for _, event := range events {
		if err := s.inner.Emit(ctx, event); err != nil {
			return err
		}
	}
	return nil
}

func isTerminalRuntimeEvent(t runtime.EventType) bool {
	return t == runtime.EventTypeComplete || t == runtime.EventTypeError || t == runtime.EventTypeCancelled
}

// sessionSink wraps a transport.Session to implement runtime.EventSink.
type sessionSink struct {
	sess       protoSender
	out        io.Writer
	nodeID     string
	sessionID  string
	background bool
	streaming  bool
	lineEnded  bool
}

func (s *sessionSink) Emit(_ context.Context, event runtime.RuntimeEvent) error {
	s.printEvent(event)
	re := &iop.RunEvent{
		RunId:      event.RunID,
		Type:       string(event.Type),
		Delta:      event.Delta,
		Message:    event.Message,
		Error:      event.Error,
		Metadata:   event.Metadata,
		Timestamp:  event.Timestamp.UnixNano(),
		SessionId:  s.sessionID,
		Background: s.background,
		NodeId:     s.nodeID,
	}
	if event.Usage != nil {
		re.Usage = &iop.Usage{
			InputTokens:  int32(event.Usage.InputTokens),
			OutputTokens: int32(event.Usage.OutputTokens),
		}
	}
	return s.sess.Send(re)
}

func (s *sessionSink) printEvent(event runtime.RuntimeEvent) {
	if s.out == nil {
		return
	}
	switch event.Type {
	case runtime.EventTypeStart:
		s.streaming = false
		s.lineEnded = true
		fmt.Fprintf(s.out, "[node-event] start run_id=%s\n", event.RunID)
	case runtime.EventTypeDelta:
		if event.Delta == "" {
			return
		}
		if !s.streaming {
			fmt.Fprint(s.out, "[node-message] ")
			s.streaming = true
		}
		fmt.Fprint(s.out, event.Delta)
		s.lineEnded = strings.HasSuffix(event.Delta, "\n")
	case runtime.EventTypeComplete:
		if s.streaming && !s.lineEnded {
			fmt.Fprintln(s.out)
		}
		s.streaming = false
		s.lineEnded = true
		fmt.Fprintf(s.out, "[node-event] complete run_id=%s detail=%q\n", event.RunID, event.Message)
	case runtime.EventTypeError:
		if s.streaming && !s.lineEnded {
			fmt.Fprintln(s.out)
		}
		s.streaming = false
		s.lineEnded = true
		fmt.Fprintf(s.out, "[node-event] error run_id=%s detail=%q\n", event.RunID, event.Error)
	case runtime.EventTypeCancelled:
		if s.streaming && !s.lineEnded {
			fmt.Fprintln(s.out)
		}
		s.streaming = false
		s.lineEnded = true
		fmt.Fprintf(s.out, "[node-event] cancelled run_id=%s\n", event.RunID)
	default:
		fmt.Fprintf(s.out, "[node-event] %s run_id=%s detail=%q\n", event.Type, event.RunID, event.Message)
	}
}

func printEdgeMessage(out io.Writer, input map[string]any) {
	if out == nil {
		return
	}
	if input == nil {
		printTaggedMessage(out, "edge-message", "")
		return
	}
	if prompt, ok := input["prompt"].(string); ok {
		printTaggedMessage(out, "edge-message", prompt)
		return
	}
	b, err := json.MarshalIndent(input, "", "  ")
	if err != nil {
		printTaggedMessage(out, "edge-message", fmt.Sprintf("%v", input))
		return
	}
	fmt.Fprintf(out, "[edge-message]\n%s\n", b)
}

func printTaggedMessage(out io.Writer, tag, message string) {
	message = strings.TrimSpace(message)
	if message == "" {
		fmt.Fprintf(out, "[%s] <empty>\n", tag)
		return
	}
	fmt.Fprintf(out, "[%s] %s\n", tag, message)
}

func structAsMap(s *structpb.Struct) map[string]any {
	if s == nil {
		return nil
	}
	return s.AsMap()
}

func sessionModeFromProto(m iop.RunSessionMode) runtime.SessionMode {
	if m == iop.RunSessionMode_RUN_SESSION_MODE_REQUIRE_EXISTING {
		return runtime.SessionModeRequireExisting
	}
	return runtime.SessionModeCreateIfMissing
}

func cancelActionFromProto(a iop.CancelAction) runtime.CancelAction {
	if a == iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION {
		return runtime.CancelActionTerminateSession
	}
	return runtime.CancelActionCancelRun
}

func normalizeSessionID(id string) string {
	if id == "" {
		return runtime.DefaultSessionID
	}
	return id
}