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

package node

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"sync"

	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/structpb"

	"iop/apps/node/internal/runtime"
	iop "iop/proto/gen/iop"
)

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
}