Cyrene_v2/internal/runtime/loop/flow.go

package loop

import (
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// graphUpdateSubject is the bus subject the FlowPublisher publishes each round's
// live flow graph on. The gateway hub subscribes to ws.> and relays it as the
// WS event type graph_update (design/19 §1).
const graphUpdateSubject = "ws.graph_update"

// allowedSessions is the predicate signature used to gate which sessions publish
// flow graphs (the owner's web sessions), so IM/stranger/sub-agent turns never
// leak into the admin panel. A nil predicate allows every session.
type allowedSessions func(sessionID string) bool

// FlowPublisher accumulates a per-session flow graph from the loop's trace spans
// and publishes ws.graph_update as each span starts/finishes, giving the frontend
// a realtime input→think→tool→output graph (design/19 §1). It is nil-tolerant
// (a nil publisher is a no-op) and safe for concurrent use; session graphs are
// serialized under one mutex because parallel tool spans finish concurrently.
//
// Ownership: one session owns one live round. A Run that declares itself a branch
// of another session (loop.WithFlowBranch) never calls StartRound/EndRound and
// publishes its spans under the owning session instead, so a sub-agent renders as
// a swimlane of the parent's graph rather than as a second graph (design/19 §2).
type FlowPublisher struct {
	mu     sync.Mutex
	publish func(subject string, payload any)
	allow  allowedSessions
	graphs map[string]*trace.RunGraph
	rounds map[string]int
}

// NewFlowPublisher builds a FlowPublisher. publish, when non-nil, is called with
// graphUpdateSubject and a *trace.FlowGraph snapshot on every span transition.
// allow, when non-nil, gates which session ids are published (returning false
// drops the session entirely). Both may be nil (a nil publish is a no-op).
func NewFlowPublisher(publish func(subject string, payload any), allow func(sessionID string) bool) *FlowPublisher {
	return &FlowPublisher{
		publish: publish,
		allow:   allow,
		graphs:  make(map[string]*trace.RunGraph),
		rounds:  make(map[string]int),
	}
}

// StartRound begins a fresh graph for sessionID, seeding an input node from the
// user text and publishing the initial frame. It resets the session's previous
// round, so each user turn is its own graph.
func (p *FlowPublisher) StartRound(sessionID, inputText string) {
	if sessionID == "" || !p.allowed(sessionID) {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	round := p.rounds[sessionID] + 1
	g := trace.NewRunGraph(sessionID)
	g.Start(round, time.Now())
	g.SeedInput(inputText)
	p.graphs[sessionID] = g
	p.rounds[sessionID] = round
	p.publishLocked(g)
}

// EndRound finalizes the round by seeding an output node from the reply and
// publishing the final frame. It keeps the graph around so a client that misses
// the final frame can replay the last snapshot on refresh.
func (p *FlowPublisher) EndRound(sessionID, outputText string) {
	if sessionID == "" || !p.allowed(sessionID) {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	g := p.graphs[sessionID]
	if g == nil {
		return
	}
	g.SeedOutput(outputText)
	p.publishLocked(g)
}

// SpanStarted records a running placeholder node for a span that just began.
// State-transition spans are omitted by the graph builder.
func (p *FlowPublisher) SpanStarted(sessionID string, sp *trace.Span) {
	if sp == nil || sessionID == "" || !p.allowed(sessionID) {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	g := p.graphs[sessionID]
	if g == nil {
		return
	}
	if g.AddSpanStart(sp.ID, sp.Name, sp.Kind, sp.ParentID, sp.StartedAt) {
		p.publishLocked(g)
	}
}

// SpanFinished folds a just-finished span into the session graph as a final node
// (status done/error) and publishes. It must be called after the span's Finish.
func (p *FlowPublisher) SpanFinished(sessionID string, sp *trace.Span) {
	if sp == nil || sessionID == "" || !p.allowed(sessionID) {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	g := p.graphs[sessionID]
	if g == nil {
		return
	}
	g.AddSpan(sp.Snapshot())
	p.publishLocked(g)
}

// publishLocked publishes a snapshot of g. Callers must hold p.mu.
func (p *FlowPublisher) publishLocked(g *trace.RunGraph) {
	if p.publish == nil {
		return
	}
	p.publish(graphUpdateSubject, g.Snapshot())
}

// allowed reports whether a session's graphs should be published.
func (p *FlowPublisher) allowed(sessionID string) bool {
	if p.allow == nil {
		return true
	}
	return p.allow(sessionID)
}

// lastUserText returns the text of the newest user message in messages, used to
// label the round's input node. It returns "" when no user text is present.
func lastUserText(messages []model.Message) string {
	for i := len(messages) - 1; i >= 0; i-- {
		if messages[i].Role != model.RoleUser {
			continue
		}
		if t := messageText(messages[i]); t != "" {
			return t
		}
	}
	return ""
}