sundynix-agentix/sundynix-dispatcher/internal/eino/compose_compiler.go

package eino

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"sync"

	"github.com/cloudwego/eino/compose"

	"github.com/sundynix/sundynix-dispatcher/internal/dsl"
	"github.com/sundynix/sundynix-shared/contract"
)

// flowSignal 是 compose 编排图的边载荷（占位）：真实数据全走 compose 本地状态(*board)，
// 边只传"该走了"的信号。注册 no-op 合并以支持 fan-in（多分支汇聚到一个节点）。
type flowSignal struct{}

var registerMergeOnce sync.Once

func registerFlowMerge() {
	registerMergeOnce.Do(func() {
		compose.RegisterValuesMergeFunc(func([]flowSignal) (flowSignal, error) { return flowSignal{}, nil })
	})
}

// executeGraph 按灰度开关选编排实现：compose.Graph（Phase C）或自研 graph.go（默认/权威）。
// 返回 (成稿, 检索来源, error)；来源供忠实度评测，两条路径都回传（compose 已对齐 graph.go）。
func (o *Orchestrator) executeGraph(ctx context.Context, t *contract.Task, tr *execTracer) (string, []string, error) {
	if composeEnabled() {
		return o.runComposeGraph(ctx, t, tr)
	}
	return o.runGraph(ctx, t, tr)
}

// runComposeGraph 把 DSL 图编译为 Eino compose.Graph 并执行（Phase C 编排归一）：
// 节点体复用现有 execDSLNode；黑板进 compose 本地状态；branch 走 AddBranch；
// DAG 触发模式让无依赖节点并行调度（效率）。编译失败即降级回自研 graph.go（安全网）。
// 返回 (成稿, 检索来源, error)：审批拒绝→errRejected、预算触顶/模型失败→fatalErr，
// 与 graph.go 终态严格对齐，否则任务会被误判 done-空。
func (o *Orchestrator) runComposeGraph(ctx context.Context, t *contract.Task, tr *execTracer) (string, []string, error) {
	return o.execComposeGraph(ctx, t, tr, nil)
}

// resumeCtx 标记一次 resume 续跑：interruptID 定位中断点，dec 是要喂给审批节点的人工决定。
type resumeCtx struct {
	interruptID string
	dec         *contract.ApprovalDecision
}

// execComposeGraph 是 compose 编排的统一入口：rc==nil 全新执行；rc!=nil 从 checkpoint resume。
// 两路共用建图/编译，仅在三处分叉：①记忆注入(仅 fresh，resume 时黑板由 checkpoint 还原)；
// ②Invoke 的 ctx(resume 注入人工决定)；③终态黑板来源(resume 时 compose 用 checkpoint 还原的
// 实例、非本闭包 b → 经 live 捕获读终态)。
func (o *Orchestrator) execComposeGraph(ctx context.Context, t *contract.Task, tr *execTracer, rc *resumeCtx) (string, []string, error) {
	registerFlowMerge()
	flow, ferr := dsl.Parse(t.Graph)
	plan := dsl.Compile(t.Graph)
	b := &board{
		uid:   meta(t, contract.MetaUserID),
		sid:   meta(t, contract.MetaSessionID),
		query: plan.Query,
	}

	// 无图/空图：退化为 compose 单轮对话。
	if ferr != nil || flow == nil || len(flow.Nodes) == 0 {
		tr.info("task", "system", "无结构化图", "按单轮对话执行（compose）")
		b.profile = o.fetchMemory(ctx, b.uid, b.query)
		b.history = o.fetchHistory(ctx, b.sid)
		o.runComposeConversation(ctx, t.ID, b, plan.System, tr, "agent")
		return b.answer, refsOf(b), b.fatalErr // 模型失败 → 上抛判 failed（对齐 graph.go）
	}

	// 邻接 + 入度（只认两端都存在的边）。
	nodeByID := make(map[string]dsl.Node, len(flow.Nodes))
	outE := make(map[string][]dsl.Edge)
	indeg := make(map[string]int, len(flow.Nodes))
	for _, n := range flow.Nodes {
		nodeByID[n.ID] = n
		indeg[n.ID] = 0
	}
	for _, e := range flow.Edges {
		if _, ok := nodeByID[e.Source]; !ok {
			continue
		}
		if _, ok := nodeByID[e.Target]; !ok {
			continue
		}
		outE[e.Source] = append(outE[e.Source], e)
		indeg[e.Target]++
	}

	// 图里无 memory 节点 → 沿用默认：注入画像+历史（与 graph.go 对齐，避免回归）。
	// 仅全新执行注入；resume 时黑板由 checkpoint 还原，重注入会覆盖已积累的执行态。
	if rc == nil {
		hasMemory := false
		for _, n := range flow.Nodes {
			if n.Kind == "memory" {
				hasMemory = true
				break
			}
		}
		if !hasMemory {
			b.profile = o.fetchMemory(ctx, b.uid, b.query)
			b.history = o.fetchHistory(ctx, b.sid)
		}
	}

	// 建 compose 图：黑板进本地状态。fresh 用本闭包 b；resume 时 compose 改用 checkpoint 还原的
	// 实例（非 b）——故节点执行时捕获 live 指针，Invoke 后据此读终态（见函数头注释③）。
	var live *board
	capture := func(bd *board) {
		if live == nil {
			live = bd // ProcessState 持状态锁 → 串行无竞态；全程同一黑板指针，设一次即可
		}
	}
	g := compose.NewGraph[flowSignal, flowSignal](
		compose.WithGenLocalState(func(context.Context) *board { return b }),
	)
	key := func(id string) string { return "n_" + id } // 节点 key 加前缀，避开 START/END 保留字

	// 1) 加节点（branch 为 passthrough，路由交给 AddBranch）。
	for _, n := range flow.Nodes {
		node := n
		if node.Kind == "branch" {
			_ = g.AddLambdaNode(key(node.ID), compose.InvokableLambda(
				func(context.Context, flowSignal) (flowSignal, error) { return flowSignal{}, nil }))
			continue
		}
		// 接了 checkpoint 后端 → 审批走中断式（compose.Interrupt 落盘释放 goroutine、抗重启），
		// 用专用 lambda（须把中断错误作为节点返回值上抛，泛型 lambda 会吞掉它）。
		// 未接后端 → 落入下方泛型 lambda 经 execDSLNode 走阻塞式 approvalNode（行为不变）。
		if node.Kind == "approval" && o.checkpoints != nil {
			_ = g.AddLambdaNode(key(node.ID), compose.InvokableLambda(o.approvalInterruptLambda(t, node, tr, capture)))
			continue
		}
		_ = g.AddLambdaNode(key(node.ID), compose.InvokableLambda(
			func(c context.Context, _ flowSignal) (flowSignal, error) {
				perr := compose.ProcessState(c, func(sc context.Context, bd *board) error {
					capture(bd) // 记下实际执行的黑板（resume 时为 checkpoint 还原的实例）
					// 上游审批拒绝 / 预算触顶 / 模型失败 → 跳过下游（对齐 graph.go 的 break 中止语义）。
					if bd.rejected || bd.fatalErr != nil {
						return nil
					}
					o.execDSLNode(sc, t, node, bd, plan, tr)
					return nil
				})
				return flowSignal{}, perr
			}))
	}

	// 2) 连边。branch 用 AddBranch（条件读 board 选下游）；其余直连；终端节点连 END。
	for _, n := range flow.Nodes {
		node := n
		outs := outE[node.ID]
		if node.Kind == "branch" {
			endNodes := map[string]bool{compose.END: true}
			for _, e := range outs {
				if _, ok := nodeByID[e.Target]; ok {
					endNodes[key(e.Target)] = true
				}
			}
			brn := node
			cond := func(c context.Context, _ flowSignal) (map[string]bool, error) {
				chosen := map[string]bool{}
				_ = compose.ProcessState(c, func(sc context.Context, bd *board) error {
					capture(bd)
					if bd.rejected || bd.fatalErr != nil {
						return nil // 已中止 → 不选任何下游，下方收口到 END
					}
					for _, tgt := range o.branchNode(brn, bd, outE[brn.ID], nodeByID, tr) {
						chosen[key(tgt)] = true
					}
					return nil
				})
				if len(chosen) == 0 {
					chosen[compose.END] = true // 没选中任何下游 → 收口到 END，避免悬挂
				}
				return chosen, nil
			}
			_ = g.AddBranch(key(node.ID), compose.NewGraphMultiBranch(cond, endNodes))
			continue
		}
		if len(outs) == 0 {
			_ = g.AddEdge(key(node.ID), compose.END)
			continue
		}
		for _, e := range outs {
			if _, ok := nodeByID[e.Target]; ok {
				_ = g.AddEdge(key(node.ID), key(e.Target))
			}
		}
	}

	// 3) 入口节点（入度 0）连 START。
	for _, n := range flow.Nodes {
		if indeg[n.ID] == 0 {
			_ = g.AddEdge(compose.START, key(n.ID))
		}
	}

	// 4) 编译（DAG 模式：无依赖节点并行调度）。接了 checkpoint 后端则挂上，使审批中断可落盘恢复。
	//    编译失败 → 降级回自研 graph.go（安全网）。
	compileOpts := []compose.GraphCompileOption{
		compose.WithNodeTriggerMode(compose.AllPredecessor),
		compose.WithGraphName("root"),
	}
	if o.checkpoints != nil {
		compileOpts = append(compileOpts, compose.WithCheckPointStore(newCheckpointStore(o.checkpoints)))
	}
	r, cerr := g.Compile(ctx, compileOpts...)
	if cerr != nil {
		tr.info("task", "system", "compose 编译失败", "退回自研 graph.go："+cerr.Error())
		return o.runGraph(ctx, t, tr) // 降级回权威实现（带 refs / 终态）
	}
	// checkpoint id = task id：审批中断时 compose 据此把整图状态（含 board）落进 store。
	var invokeOpts []compose.Option
	if o.checkpoints != nil {
		invokeOpts = append(invokeOpts, compose.WithCheckPointID(t.ID))
	}
	// resume 续跑：把人工决定注入 ctx，喂给中断点的审批节点（fresh 则原样 invoke）。
	invokeCtx := ctx
	if rc != nil {
		invokeCtx = compose.ResumeWithData(ctx, rc.interruptID, rc.dec)
	}
	if _, ierr := r.Invoke(invokeCtx, flowSignal{}, invokeOpts...); ierr != nil {
		// HITL 审批中断：checkpoint 已落、任务停在 waiting（审批 lambda 内已置）→ 落 resume 记录
		// 供决定到达时续跑，上抛哨兵让 Handle 释放 goroutine 不收尾。
		if info, ok := compose.ExtractInterruptInfo(ierr); ok {
			id := firstInterruptID(info)
			o.persistResume(ctx, t, id, tr)
			tr.info("task", "approval", "已中断等待审批", "checkpoint 已落，释放执行；interrupt="+id)
			return b.answer, nil, errInterrupted
		}
		tr.info("task", "system", "compose 执行告警", ierr.Error()) // 非中断告警：副作用已落 board，按终态收尾
	}

	// 终态读「实际执行的黑板」：fresh=闭包 b；resume=compose 用 checkpoint 还原的实例(live)。
	fb := live
	if fb == nil {
		fb = b
	}

	// 终态对齐 graph.go：审批拒绝 / 预算触顶 / 模型失败要显式上抛，否则任务误判 done-空。
	if fb.rejected {
		o.clearResume(ctx, t.ID)           // 拒绝是终态，清 resume 记录 + checkpoint
		return fb.answer, nil, errRejected // 合法终态，Handle 据此判 rejected 并优雅收尾
	}
	if fb.fatalErr != nil {
		return fb.answer, nil, fb.fatalErr // 上抛 → Handle 判 failed（带原因）
	}

	// 图里无 agent 节点（纯工具/检索图）也要出一段答复。
	if fb.answer == "" {
		o.runComposeConversation(ctx, t.ID, fb, plan.System, tr, "agent")
		if fb.fatalErr != nil { // 兜底对话也可能触预算顶 / 模型失败
			return fb.answer, nil, fb.fatalErr
		}
	}
	o.clearResume(ctx, t.ID)          // 成功收尾：清 resume 记录 + checkpoint
	return fb.answer, refsOf(fb), nil // 成功：带回检索来源供忠实度评测
}

// execDSLNode 执行一个非 branch 的 DSL 节点（compose 编译器用；节点体与 graph.go 一致，
// 区别仅 agent 直接走 runAgent/runReactAgent——compose 已是编排层，不再二次套 compose）。
func (o *Orchestrator) execDSLNode(ctx context.Context, t *contract.Task, n dsl.Node, b *board, plan dsl.Plan, tr *execTracer) {
	switch n.Kind {
	case "input":
		if txt := cstr(n.Config, "text"); txt != "" {
			b.query = txt
		}
		tr.info("input:"+n.ID, "system", labelOf(n, "输入"), truncate(b.query, 80))
	case "memory":
		if cbool(n.Config, "profile") {
			b.profile = o.fetchMemory(ctx, b.uid, b.query)
		}
		if cbool(n.Config, "history") {
			b.history = o.fetchHistory(ctx, b.sid)
		}
		tr.info("memory:"+n.ID, "memory", labelOf(n, "记忆"),
			fmt.Sprintf("画像 %d 字 · 历史 %d 条", len([]rune(b.profile)), len(b.history)))
	case "retriever":
		o.retrieverNode(ctx, n, b, tr)
	case "tool":
		o.execToolNode(ctx, t.ID, n, b, tr)
	case "agent":
		sys := firstNonEmpty(cstr(n.Config, "system"), plan.System)
		if cbool(n.Config, "autonomous") {
			o.runReactAgent(ctx, t.ID, b, sys, n, tr, "agent:"+n.ID)
		} else {
			o.runAgent(ctx, t.ID, b, sys, tr, "agent:"+n.ID)
		}
	case "coordinator": // 多智能体协调：orchestrator 自主把子任务派给专家(agent-as-tool)再综合
		o.runCoordinator(ctx, t.ID, b, firstNonEmpty(cstr(n.Config, "system"), plan.System), n, tr, "coordinator:"+n.ID)
	case "aggregate":
		merged := aggregate(cstr(n.Config, "strategy"), append(append([]string{}, b.refs...), b.toolOut...))
		b.refs, b.toolOut = merged, nil
		tr.info("aggregate:"+n.ID, "system", labelOf(n, "汇聚"), "策略："+firstNonEmpty(cstr(n.Config, "strategy"), "拼接"))
	case "approval":
		// 阻塞式审批（未接 checkpoint 后端时走这里）：阻塞等批准/拒绝，拒绝/超时置 b.rejected
		// （节点入口守卫据此中止下游）。返回的放行清单由静态边 + rejected 守卫接管，丢弃。
		// 接了 checkpoint 后端时，审批已在建图阶段改走 approvalInterruptLambda，不会到这。
		o.approvalNode(ctx, t.ID, n, b, tr, nil)
	case "render":
		o.renderNode(ctx, t.ID, n, b, tr)
	case "map":
		o.mapNode(ctx, t.ID, n, b, tr)
	case "output":
		tr.info("output:"+n.ID, "system", labelOf(n, "输出"), "目标："+firstNonEmpty(cstr(n.Config, "target"), "屏幕"))
	default:
		tr.info(n.Kind+":"+n.ID, "system", labelOf(n, n.Kind), "未识别节点，跳过")
	}
}

// approvalInterruptLambda 是中断式 HITL 审批节点体（compose checkpoint）：
//   - 首次执行 → 发待审事件 + 置 waiting + compose.Interrupt：compose 把整图状态（含 board，
//     经 *board 的自定义 JSON 序列化）落进 checkpoint store 并返回中断错误，
//     runComposeGraph 据 ExtractInterruptInfo 上抛 errInterrupted、Handle 释放 goroutine；
//   - resume 流 → GetResumeContext 取人工决定，落黑板（批准放行 / 拒绝置 rejected 中止下游）。
//
// 上游已中止（rejected/fatalErr）则直接跳过，不发待审、不中断。
func (o *Orchestrator) approvalInterruptLambda(t *contract.Task, n dsl.Node, tr *execTracer, capture func(*board)) func(context.Context, flowSignal) (flowSignal, error) {
	return func(ctx context.Context, _ flowSignal) (flowSignal, error) {
		// resume 流：应用人工决定。
		if isResume, hasData, dec := compose.GetResumeContext[*contract.ApprovalDecision](ctx); isResume {
			approved := hasData && dec != nil && dec.Approved
			perr := compose.ProcessState(ctx, func(_ context.Context, b *board) error {
				capture(b) // resume 时这是 checkpoint 还原的黑板，须捕获以读终态
				o.applyApprovalDecision(t.ID, n, b, dec, approved, tr)
				return nil
			})
			return flowSignal{}, perr
		}
		// 首次执行：读待审摘要；上游已中止则跳过（不发待审、不中断）。
		var title, summary string
		skip := false
		_ = compose.ProcessState(ctx, func(_ context.Context, b *board) error {
			capture(b)
			if b.rejected || b.fatalErr != nil {
				skip = true
				return nil
			}
			title, summary = approvalSummary(n, b)
			return nil
		})
		if skip {
			return flowSignal{}, nil
		}
		// 发待审 + 置 waiting，然后中断（compose 落 checkpoint、释放 goroutine）。
		tr.emit("approval:"+n.ID, "approval", "await", title, summary, 0)
		o.setStatus(t.ID, contract.TaskWaiting, title)
		return flowSignal{}, compose.Interrupt(ctx, title)
	}
}

// ---- HITL resume：中断记录的持久化与续跑 ----

// errNoPending 表示无 resume 记录（任务未中断 / 记录已过 TTL 被清）。
var errNoPending = errors.New("no pending approval record")

// pendingApproval 是审批中断的 resume 记录：决定到达时据此重建任务并从 checkpoint 续跑。
// 与 compose checkpoint 同桶分键存（checkpoint 键=task_id，本记录键=pending:task_id）。
// 存 Task 原文是为了让决定能在另一 goroutine / 重启后的新进程里独立重建任务再 resume。
type pendingApproval struct {
	InterruptID string          `json:"interrupt_id"`
	Task        json.RawMessage `json:"task"`
}

// pendingKey 是 resume 记录的 KV 键。注意：NATS JetStream KV 键只允许 [-/_=.a-zA-Z0-9]，
// 冒号等字符会被拒（nats: invalid key）——故用下划线分隔，不可改回冒号。
func pendingKey(taskID string) string { return "pending_" + taskID }

// firstInterruptID 取本次中断的首个 interrupt id（审批是单点中断，取首个即可）。
func firstInterruptID(info *compose.InterruptInfo) string {
	if info != nil && len(info.InterruptContexts) > 0 {
		return info.InterruptContexts[0].ID
	}
	return ""
}

// persistResume 落 resume 记录（无 checkpoint 后端则跳过；失败仅告警不阻断——checkpoint 仍在，
// 大不了靠人工/兜底重投）。
func (o *Orchestrator) persistResume(ctx context.Context, t *contract.Task, interruptID string, tr *execTracer) {
	if o.checkpoints == nil {
		return
	}
	tb, err := t.Marshal()
	if err != nil {
		tr.info("task", "system", "resume 记录序列化失败", err.Error())
		return
	}
	rec, _ := json.Marshal(pendingApproval{InterruptID: interruptID, Task: tb})
	if err := o.checkpoints.Put(ctx, pendingKey(t.ID), rec); err != nil {
		// 关键：落盘失败则决定到达时无从 resume，任务永卡 waiting → 大声 log（非仅 exec 轨迹）。
		log.Printf("[eino] resume 记录落盘失败 task=%s: %v（该任务将无法恢复）", t.ID, err)
		tr.info("task", "system", "resume 记录落盘失败", err.Error())
	}
}

// loadResume 取 resume 记录（决定消费者据 task_id 续跑用）。无记录返回 errNoPending。
func (o *Orchestrator) loadResume(ctx context.Context, taskID string) (*pendingApproval, error) {
	if o.checkpoints == nil {
		return nil, errNoPending
	}
	data, ok, err := o.checkpoints.Get(ctx, pendingKey(taskID))
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, errNoPending
	}
	var p pendingApproval
	if err := json.Unmarshal(data, &p); err != nil {
		return nil, err
	}
	return &p, nil
}

// clearResume 终态清理：删 resume 记录 + compose checkpoint（幂等；无后端则跳过）。
func (o *Orchestrator) clearResume(ctx context.Context, taskID string) {
	if o.checkpoints == nil {
		return
	}
	_ = o.checkpoints.Delete(ctx, pendingKey(taskID))
	_ = o.checkpoints.Delete(ctx, taskID)
}

// ResumeApproval 续跑一个等待审批的任务：据 task_id 取 resume 记录定位中断点，把人工决定喂给
// 审批节点、从 checkpoint 重入图。供决定消费者调用（增量3b）。返回终态语义同 runComposeGraph：
// 续跑中再遇审批 → 仍返回 errInterrupted（重新落记录）；批准跑完 → 成稿+refs；拒绝 → errRejected。
func (o *Orchestrator) ResumeApproval(ctx context.Context, t *contract.Task, dec *contract.ApprovalDecision, tr *execTracer) (string, []string, error) {
	p, err := o.loadResume(ctx, t.ID)
	if err != nil {
		return "", nil, err
	}
	return o.execComposeGraph(ctx, t, tr, &resumeCtx{interruptID: p.InterruptID, dec: dec})
}

// HandleApprovalDecision 是审批决定持久消费者的回调：据 task_id 取 resume 记录续跑并收尾。
// 无 resume 记录（阻塞态任务的决定 / 已恢复 / 已过期）则忽略——本消费者只管中断/恢复模型。
// 幂等：决定重投时 resume 记录已被成功收尾清掉 → loadResume 落空 → 安全跳过。
func (o *Orchestrator) HandleApprovalDecision(ctx context.Context, dec *contract.ApprovalDecision) {
	if o.checkpoints == nil || dec == nil || dec.TaskID == "" {
		return
	}
	p, err := o.loadResume(ctx, dec.TaskID)
	if errors.Is(err, errNoPending) {
		return // 非中断态任务（阻塞模型自己经 WaitApproval 收）/ 已处理过
	}
	if err != nil {
		log.Printf("[eino] 取 resume 记录失败 task=%s: %v", dec.TaskID, err)
		return
	}
	var t contract.Task
	if uerr := json.Unmarshal(p.Task, &t); uerr != nil || t.ID == "" {
		log.Printf("[eino] resume 记录损坏 task=%s: %v", dec.TaskID, uerr)
		o.clearResume(ctx, dec.TaskID) // 坏记录清掉，避免决定无限重投
		return
	}
	tr := o.tracer(t.ID)
	defer tr.done()
	answer, refs, rerr := o.ResumeApproval(ctx, &t, dec, tr)
	o.finishResumed(ctx, &t, answer, refs, rerr)
}

// finishResumed 给 resume 续跑收尾，语义对齐 Handle 尾段（中断/拒绝/预算/失败/成功）。
func (o *Orchestrator) finishResumed(ctx context.Context, t *contract.Task, answer string, refs []string, err error) {
	switch {
	case errors.Is(err, errInterrupted):
		return // 续跑中又遇审批：保持 waiting，等下一个决定（resume 记录已重新落盘）
	case errors.Is(err, errRejected):
		if answer != "" {
			_ = o.sink.PublishToken(t.ID, []byte(answer))
		}
		_ = o.sink.CompleteStream(t.ID)
		o.breaker.Report(true) // 拒绝是人为决策，非后端故障
		o.setStatus(t.ID, contract.TaskRejected, truncate(answer, 120))
		return
	case errors.Is(err, errBudget):
		if answer != "" {
			_ = o.sink.PublishToken(t.ID, []byte(answer))
		}
		_ = o.sink.PublishToken(t.ID, []byte("\n\n⚠️ 已达单任务 token 预算上限，自动中止。"))
		_ = o.sink.CompleteStream(t.ID)
		o.breaker.Report(true)
		o.setStatus(t.ID, contract.TaskFailed, "token 预算超限")
		return
	case err != nil:
		_ = o.sink.CompleteStream(t.ID)
		o.breaker.Report(false)
		o.finishStatus(t.ID, err)
		return
	}
	// 成功：收尾流 + 判 done + 评测（低分纠偏）+ 落历史（同 Handle）。
	_ = o.sink.CompleteStream(t.ID)
	o.breaker.Report(true)
	o.finishStatus(t.ID, nil)
	go func() {
		query := dsl.Compile(t.Graph).Query
		final := o.evaluate(t, query, answer, refs)
		o.memorize(t, final)
	}()
}