1086 lines
33 KiB
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
|
|
}
|