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

package node

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

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

	runtime "iop/packages/go/execution"
)

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

type noopSender struct{}

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

// 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)
	event.Metadata = cloneLivenessMetadata(event.Metadata)
	return s.sess.Send(runEventToProto(event, s.nodeID, s.sessionID, s.background))
}

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 normalizeSessionID(id string) string {
	if id == "" {
		return runtime.DefaultSessionID
	}
	return id
}