Cyrene_v2/internal/runtime/loop/flow.go

package loop

import (
	"log/slog"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"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"

// flowLogSource 是发布器日志的 source 标签（落盘失败记 WARN；AGENTS §10 禁裸日志）。
const flowLogSource = "runtime/loop/flow"

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

// RoundRecorder 把一条**已收尾轮次**的图落盘（design/19 §4 的 P3「多轮历史回放」；
// 切片 F2）。生产实现是 `internal/runtime/flowgraph` 的 `*SQLStore`。
//
// 语义约定（实现方必须遵守）：**每轮一条**（只在 EndRound 里调用，绝不在每个 span
// 上调用），失败**上抛**给发布器（发布器记 WARN 后继续）—— 图是可视化资产，落盘失败
// 绝不允许影响会话收尾，也绝不改变实时帧。
//
// 白名单不在实现方：EndRound 对发布会话范围外的会话**直接返回**，因此落盘出口天然
// 只被白名单内的会话触发（复用 flowSessionAllowed，不存在第二份白名单）。
type RoundRecorder interface {
	RecordRound(g *trace.FlowGraph) error
}

// FlowPublisherOption 配置 FlowPublisher 的可选行为（落盘出口 + 日志）。
type FlowPublisherOption func(*FlowPublisher)

// WithRoundRecorder 挂一个「每轮收尾时落一条」的落盘出口。传 nil 等价于不落盘。
func WithRoundRecorder(rec RoundRecorder) FlowPublisherOption {
	return func(p *FlowPublisher) { p.record = rec }
}

// WithFlowLogger 给发布器一个 logger：落盘失败要**可见**（记 WARN），但绝不把图内容
// 写进日志（图里可能含用户输入文本，只记 session/turn 标识）。
func WithFlowLogger(log *slog.Logger) FlowPublisherOption {
	return func(p *FlowPublisher) { p.log = log }
}

// 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
	// record 是每轮收尾的落盘出口（design/19 §4 P3；切片 F2）。nil = 不落盘
	// （保持 F1 之前的纯内存行为，且测试可只测发布）。
	record RoundRecorder
	// log 用于把落盘失败记成 WARN（带 source）。nil 时静默（测试/未接线）。
	log *slog.Logger
}

// NewFlowPublisher 构建一个 FlowPublisher。publish 非 nil 时，会在
// 每次 span 状态迁移时以 graphUpdateSubject 与 *trace.FlowGraph
// 快照调用。allow 非 nil 时，控制哪些 session id 可发布（返回
// false 则整个会话被丢弃）。两者都可以为 nil（nil publish 为空操作）。
//
// 可变项经 Functional Options 传入（不改变既有两参调用）：落盘出口（F2）与日志。
func NewFlowPublisher(publish func(subject string, payload any), allow func(sessionID string) bool, opts ...FlowPublisherOption) *FlowPublisher {
	p := &FlowPublisher{
		publish: publish,
		allow:   allow,
		graphs:  make(map[string]*trace.RunGraph),
		rounds:  make(map[string]int),
	}
	for _, opt := range opts {
		if opt != nil {
			opt(p)
		}
	}
	return p
}

// 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 节点并发布最终帧，
// 以此收尾该轮次。它保留该图，使错过最终帧的客户端能重放最后一份快照 ——
// **重放的出口是 `LastGraph`**（网关只读查询，契约 v0.71），而不是又发一帧：
// 实时帧向来只推给当时在线的连接，刷新/晚开的页面只能靠只读拉取补齐
// （`graph_update` 本身不重放）。
//
// 收尾也是**落盘**的唯一时机（切片 F2；design/19 §4 的 P3「多轮历史回放」）：每轮落
// **一条**（绝不是每 span 一条），且只在该会话处于发布会话白名单内时发生（本方法对
// 白名单外的会话直接返回 ⇒ 落盘出口天然只被白名单会话触发）。深拷贝与写库在**放锁
// 之后**做：一次库写绝不把其他会话的 span 发布挡住；失败只记 WARN，绝不影响本轮收尾。
func (p *FlowPublisher) EndRound(sessionID, outputText string) {
	if sessionID == "" || !p.allowed(sessionID) {
		return
	}
	p.mu.Lock()
	g := p.graphs[sessionID]
	if g == nil {
		p.mu.Unlock()
		return
	}
	g.SeedOutput(outputText)
	p.publishLocked(g)
	var snap *trace.FlowGraph
	if p.record != nil {
		snap = g.Snapshot().Clone()
	}
	p.mu.Unlock()

	if snap != nil {
		if err := p.record.RecordRound(snap); err != nil {
			p.logPersistError(sessionID, snap.ID, err)
		}
	}
}

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

// LastGraph 返回 sessionID **最近一轮**链路图的一份只读深拷贝快照，供
// 管理面板在「轮次已经结束、实时帧已经错过」时补拉（契约 v0.71 的
// `GET /api/v1/admin/flow/graph`；design/19 §4 的 P3「多轮历史回放」的
// 最小一步）。
//
// 它是**只读**的：不发布、不改图、不动轮次计数 —— 因此补拉绝不会在
// 管理面板里凭空造出一帧 `graph_update`，也绝不改任何会话的实时图。
//
// 该会话还没有图时返回 nil：从未跑过带 span 的轮次、或该会话本就**不在
// 发布会话范围内**（`StartRound` 对白名单外的会话直接返回 ⇒ 从不建图，
// 隔离渠道会话因此天然查不到）。调用方**必须**把 nil 读作「确实没有图」，
// 绝不允许据此补画任何节点或假图。
//
// 返回的是深拷贝（`trace.FlowGraph.Clone`）：读者在锁外编码时，
// 活图可能正被下一帧 span 就地合并，共享底层 map/切片会撕裂。
func (p *FlowPublisher) LastGraph(sessionID string) *trace.FlowGraph {
	if p == nil || sessionID == "" {
		return nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	g := p.graphs[sessionID]
	if g == nil {
		return nil
	}
	return g.Snapshot().Clone()
}

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

// logPersistError 把一次落盘失败记成带来源的 WARN。它只记 session/turn 标识，
// **绝不打印图内容**（图里可能含用户输入文本）；logger 未接线时静默。
func (p *FlowPublisher) logPersistError(sessionID, turnID string, err error) {
	if p.log == nil {
		return
	}
	logger.WithSource(p.log, flowLogSource).Warn("flow round persist failed",
		"session", sessionID, "turn", turnID, "err", err)
}

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