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 是 FlowPublisher 发布每轮实时
// flow 图所用的 bus subject。gateway hub 订阅 ws.> 并将其
// 作为 WS 事件类型 graph_update 转发（design/19 §1）。
const graphUpdateSubject = "ws.graph_update"

// allowedSessions 是用于控制哪些会话可发布
// flow 图的谓词签名（主人的 web 会话），使 IM/陌生人/子代理轮次绝不
// 泄露进管理面板。谓词为 nil 时允许所有会话。
type allowedSessions func(sessionID string) bool

// FlowPublisher 从循环的 trace span 累积每会话 flow 图，
// 并在每个 span 开始/结束时发布 ws.graph_update，给前端
// 一张实时的 input→think→tool→output 图（design/19 §1）。它容忍 nil
// （nil 发布者为空操作）且并发安全；会话图
// 在同一把互斥锁下串行化，因为并行的工具 span 会并发完成。
//
// 归属：一个会话拥有一个实时轮次。声明自己是另一会话分支的
// Run（loop.WithFlowBranch）绝不调用 StartRound/EndRound，
// 而是把其 span 发布到所属会话下，因此子代理渲染为
// 父图的一条泳道，而不是第二张图（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 构建一个 FlowPublisher。publish 非 nil 时，会在
// 每次 span 状态迁移时以 graphUpdateSubject 与 *trace.FlowGraph
// 快照调用。allow 非 nil 时，控制哪些 session id 可发布（返回
// false 则整个会话被丢弃）。两者都可以为 nil（nil publish 为空操作）。
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 为 sessionID 开始一张新图，用用户文本
// 播种一个 input 节点并发布初始帧。它重置该会话的上一
// 轮次，因此每个用户 turn 都是自己的一张图。
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 用回复播种一个 output 节点并发布最终帧，
// 以此收尾该轮次。它保留该图，使错过最终帧的客户端
// 刷新时能重放最后一份快照。
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)
}

// MarkInputSteered 给本轮的 input 节点打上 `steered:true` 标记并发布一帧
// （design/15 §5 / 裁决 13.4.5「注入时给 input 节点打 steered:true 标记」）。
// 它让链路图如实表达「这一轮的输入在回合中途被追加过」——图仍然只有一条主干
// （I1：同会话恒只有一个在飞回合），而不是被画成第二轮。
//
// nil 接收者、空会话、发布白名单外的会话、或该会话还没有图时都是空操作；标记
// 只并入既有 input 节点（Upsert 的合并语义），绝不新建节点、绝不改动它的标签。
func (p *FlowPublisher) MarkInputSteered(sessionID string) {
	if p == nil || sessionID == "" || !p.allowed(sessionID) {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	g := p.graphs[sessionID]
	if g == nil {
		return
	}
	g.Upsert(trace.NodeSpec{
		ID:    "input",
		Kind:  "input",
		Props: map[string]any{"steered": true},
	})
	p.publishLocked(g)
}

// SpanStarted 为一个刚开始的 span 记录一个运行中的占位
// 节点。state_transition span 会被图构建器省略。
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 把一个刚结束的 span 作为最终节点（状态 done/error）
// 并入会话图并发布。必须在 span 的 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 发布 g 的一份快照。调用方必须持有 p.mu。
func (p *FlowPublisher) publishLocked(g *trace.RunGraph) {
	if p.publish == nil {
		return
	}
	p.publish(graphUpdateSubject, g.Snapshot())
}

// allowed 报告某会话的图是否应当发布。
func (p *FlowPublisher) allowed(sessionID string) bool {
	if p.allow == nil {
		return true
	}
	return p.allow(sessionID)
}

// lastUserText 返回 messages 中最新用户消息的文本，用于
// 标注该轮的 input 节点。没有用户文本时返回 ""。
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 ""
}