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"
	"strings"
	"time"

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

	"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
	out    io.Writer
	logger *zap.Logger
}

// New creates a Node. It satisfies transport.Handler.
// out receives console debug output; pass os.Stdout for production, io.Discard in tests.
func New(
	nodeID string,
	router runtime.Router,
	st *store.Store,
	out io.Writer,
	logger *zap.Logger,
) *Node {
	if out == nil {
		out = os.Stdout
	}
	return &Node{
		nodeID: nodeID,
		router: router,
		store:  st,
		runs:   newRunManager(),
		out:    out,
		logger: logger,
	}
}

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

	spec, adapter, err := n.router.ResolveAdapter(ctx, rr)
	if err != nil {
		return fmt.Errorf("node: resolve: %w", err)
	}

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

	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)

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

	run := func() error {
		defer cancel()
		defer n.runs.deregister(spec.RunID)
		defer close(h.done)

		execErr := adapter.Execute(execCtx, spec, sink)
		n.completeRun(spec, execErr)
		return execErr
	}

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

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, ok := n.router.GetAdapter(req.GetAdapter())
		if !ok {
			return fmt.Errorf("node: adapter %q not found", req.GetAdapter())
		}
		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
	}
}

// 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()),
	)

	adapter, ok := n.router.GetAdapter(req.GetAdapter())
	if !ok {
		return &iop.NodeCommandResponse{
			RequestId: req.GetRequestId(),
			Type:      req.GetType(),
			Adapter:   req.GetAdapter(),
			Target:    req.GetTarget(),
			SessionId: req.GetSessionId(),
			Error:     fmt.Sprintf("node: adapter %q not found", req.GetAdapter()),
		}, nil
	}

	handler, ok := adapter.(runtime.CommandHandler)
	if !ok {
		return &iop.NodeCommandResponse{
			RequestId: req.GetRequestId(),
			Type:      req.GetType(),
			Adapter:   req.GetAdapter(),
			Target:    req.GetTarget(),
			SessionId: req.GetSessionId(),
			Error:     fmt.Sprintf("node: adapter %q does not support commands", req.GetAdapter()),
		}, nil
	}

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

	domainReq := toDomainCommandRequest(req)
	domainResp, err := handler.HandleCommand(execCtx, domainReq)
	if err != nil {
		return &iop.NodeCommandResponse{
			RequestId: req.GetRequestId(),
			Type:      req.GetType(),
			Adapter:   req.GetAdapter(),
			Target:    req.GetTarget(),
			SessionId: req.GetSessionId(),
			Error:     err.Error(),
		}, nil
	}

	return toProtoCommandResponse(domainResp), nil
}

func toDomainCommandRequest(req *iop.NodeCommandRequest) runtime.CommandRequest {
	var cmdType runtime.CommandType
	if req.GetType() == iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS {
		cmdType = runtime.CommandTypeUsageStatus
	}
	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 {
	var cmdType iop.NodeCommandType
	if resp.Type == runtime.CommandTypeUsageStatus {
		cmdType = iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS
	}
	out := &iop.NodeCommandResponse{
		RequestId: resp.RequestID,
		Type:      cmdType,
		Adapter:   resp.Adapter,
		Target:    resp.Target,
		SessionId: resp.SessionID,
	}
	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
}

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