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

1086 lines
33 KiB
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 (retained for compatibility; not used for admission)
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 retained as a compatibility argument but is no longer
// used for admission. Node-wide concurrency limits have been removed; per-
// adapter MaxConcurrency from Capabilities is the sole admission gate.
// 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 {
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)
// 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()
if !runSink.hasTerminalObserved() {
if synthErr := n.synthAndEmitTerminal(ctx, runSink, spec, execErr); synthErr != nil {
if execErr == nil {
execErr = synthErr
}
}
}
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()
}
// OnProviderTunnelRequest handles an incoming ProviderTunnelRequest from a transport Session.
func (n *Node) OnProviderTunnelRequest(ctx context.Context, sess *transport.Session, req *iop.ProviderTunnelRequest) error {
n.logger.Info("provider tunnel request received",
zap.String("run_id", req.GetRunId()),
zap.String("tunnel_id", req.GetTunnelId()),
zap.String("adapter", req.GetAdapter()),
zap.String("target", req.GetTarget()),
)
tr := runtime.ProviderTunnelRequest{
RunID: req.GetRunId(),
TunnelID: req.GetTunnelId(),
Adapter: req.GetAdapter(),
Target: req.GetTarget(),
Method: req.GetMethod(),
Path: req.GetPath(),
Headers: req.GetHeaders(),
Body: req.GetBody(),
Stream: req.GetStream(),
TimeoutSec: int(req.GetTimeoutSec()),
Metadata: req.GetMetadata(),
SessionID: req.GetSessionId(),
}
n.configSetMu.RLock()
adapter, err := n.router.LookupAdapter(tr.Adapter)
n.configSetMu.RUnlock()
if err != nil {
n.sendTunnelError(sess, tr, fmt.Errorf("node: lookup adapter: %w", err))
return fmt.Errorf("node: lookup adapter: %w", err)
}
tunnelAdapter, ok := adapter.(runtime.ProviderTunnelAdapter)
if !ok {
err := fmt.Errorf("node: adapter %q does not support tunneling", tr.Adapter)
n.sendTunnelError(sess, tr, err)
return err
}
var sender protoSender = noopSender{}
nodeID := n.nodeID
nodeAlias := ""
if sess != nil {
if sess.IsAlive() {
sender = sess
}
nodeID = sess.NodeID()
nodeAlias = sess.Alias()
}
sink := &tunnelSink{
sess: sender,
nodeID: nodeID,
nodeAlias: nodeAlias,
}
execCtx, cancel := context.WithCancel(ctx)
if tr.TimeoutSec > 0 {
execCtx, cancel = context.WithTimeout(ctx, time.Duration(tr.TimeoutSec)*time.Second)
}
defer cancel()
h := &runHandle{
runID: tr.RunID,
adapter: tr.Adapter,
target: tr.Target,
sessionID: normalizeSessionID(tr.SessionID),
cancel: cancel,
done: make(chan struct{}),
}
n.runs.register(h)
defer n.runs.deregister(tr.RunID)
defer close(h.done)
if err := tunnelAdapter.TunnelProvider(execCtx, tr, sink); err != nil {
n.logger.Warn("provider tunnel error",
zap.String("run_id", tr.RunID),
zap.String("tunnel_id", tr.TunnelID),
zap.Error(err),
)
return err
}
return nil
}
func (n *Node) sendTunnelError(sess *transport.Session, tr runtime.ProviderTunnelRequest, err error) {
if sess == nil || !sess.IsAlive() {
return
}
tf := &iop.ProviderTunnelFrame{
RunId: tr.RunID,
TunnelId: tr.TunnelID,
Sequence: 0,
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
Error: err.Error(),
Timestamp: time.Now().UnixNano(),
NodeId: sess.NodeID(),
NodeAlias: sess.Alias(),
}
_ = sess.Send(tf)
}
type tunnelSink struct {
sess protoSender
nodeID string
nodeAlias string
}
func (s *tunnelSink) EmitTunnelFrame(ctx context.Context, frame runtime.ProviderTunnelFrame) error {
var usage *iop.Usage
if frame.Usage != nil {
usage = &iop.Usage{
InputTokens: int32(frame.Usage.InputTokens),
OutputTokens: int32(frame.Usage.OutputTokens),
ReasoningTokens: int32(frame.Usage.ReasoningTokens),
CachedInputTokens: int32(frame.Usage.CachedInputTokens),
}
}
protoKind := iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_UNSPECIFIED
switch frame.Kind {
case runtime.ProviderTunnelFrameKindResponseStart:
protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START
case runtime.ProviderTunnelFrameKindBody:
protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY
case runtime.ProviderTunnelFrameKindEnd:
protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END
case runtime.ProviderTunnelFrameKindError:
protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR
case runtime.ProviderTunnelFrameKindUsage:
protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE
}
tf := &iop.ProviderTunnelFrame{
RunId: frame.RunID,
TunnelId: frame.TunnelID,
Sequence: frame.Sequence,
Kind: protoKind,
StatusCode: int32(frame.StatusCode),
Headers: frame.Headers,
Body: frame.Body,
End: frame.End,
Error: frame.Error,
Usage: usage,
Metadata: frame.Metadata,
Timestamp: frame.Timestamp.UnixNano(),
NodeId: s.nodeID,
NodeAlias: s.nodeAlias,
}
if s.sess != nil {
return s.sess.Send(tf)
}
return nil
}
// 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()
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
}
// 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 concurrency is no longer applied live to admission gates.
// Config metadata is preserved in currentConfigSet.Runtime.
// Per-adapter capacity from Capabilities().MaxConcurrency is the sole admission source.
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
}
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()
// 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
terminalObserved bool
deferred []runtime.RuntimeEvent
}
func (s *terminalDeferringSink) Emit(ctx context.Context, event runtime.RuntimeEvent) error {
s.mu.Lock()
if isTerminalRuntimeEvent(event.Type) {
s.terminalObserved = true
}
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 (s *terminalDeferringSink) hasTerminalObserved() bool {
s.mu.Lock()
defer s.mu.Unlock()
return s.terminalObserved
}
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),
ReasoningTokens: int32(event.Usage.ReasoningTokens),
CachedInputTokens: int32(event.Usage.CachedInputTokens),
}
}
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.EventTypeReasoningDelta:
if event.Delta == "" {
return
}
if s.streaming && !s.lineEnded {
fmt.Fprintln(s.out)
}
s.streaming = false
s.lineEnded = true
printPrefixedLines(s.out, "[node-reasoning] ", event.Delta)
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 printPrefixedLines(out io.Writer, prefix, text string) {
text = strings.TrimRight(text, "\n")
if text == "" {
return
}
for _, line := range strings.Split(text, "\n") {
fmt.Fprintf(out, "%s%s\n", prefix, line)
}
}
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
}