ai-agent/workflow/service/flow/exec_hub.go

package flow

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	wsCommon "gitea.redpowerfuture.com/red-future/common/websocket"
)

// ====================== 执行事件中枢（attach to running execution） ======================
//
// execHub 是单次工作流执行的事件中枢：把节点进度（node_start/node_complete）与终态
// （flow_complete/error）广播到所有订阅的 WS 连接，并暴露执行取消入口（CancelByUser）。
//
// 背景：恢复例程捞起的执行没有 WS 连接、进度被丢弃；用户执行中再点"执行"时，现有代码只发
// round_start "正在执行中" 就返回，收不到进度也无法取消。execHub 让任意时刻建立的连接都能
// 订阅到同一 session+flow 运行中执行的后续进度，并通过 workflow_cancel 停止它。
//
// 生命周期：执行方（handleExecute 的新执行/断点续跑、recoverExecution 的恢复）创建并接管
// （SetCancel + MarkOwned）hub，终态落库后 Publish 终态消息并 Close（退订全部连接、注销）。
// 候选 hub 由 handleExecute 提前注册（registerHubIfAbsent），同键已有运行中执行则复用其 hub。

type execHub struct {
	sessionId string
	flowId    int64
	execId    int64 // 已知后设置，仅用于日志

	mu     sync.Mutex
	subs   map[*wsCommon.WsConnection]struct{}
	cancel context.CancelFunc // 取消执行（用户执行=execCancel / 恢复=topCancel）
	owned  bool               // 已被执行方接管（MarkOwned）；未接管视为候选/占位
	closed bool

	closeOnce     sync.Once
	done          chan struct{}
	userCancelled atomic.Bool
}

func newExecHub(sessionId string, flowId int64) *execHub {
	return &execHub{
		sessionId: sessionId,
		flowId:    flowId,
		subs:      make(map[*wsCommon.WsConnection]struct{}),
		done:      make(chan struct{}),
	}
}

// ====================== 注册表（按 sessionId+flowId） ======================

var (
	hubRegMu sync.Mutex
	hubReg   = make(map[string]*execHub)
)

func hubKey(sessionId string, flowId int64) string {
	return fmt.Sprintf("%s\x00%d", sessionId, flowId)
}

// registerHubIfAbsent 注册 hub；同键已有则返回现有 hub（候选丢弃，返回 created=false）。
// 保证同一进程内同 session+flow 同时只有一个 hub 对象，后续连接统一订阅到它。
func registerHubIfAbsent(sessionId string, flowId int64, hub *execHub) (existing *execHub, created bool) {
	key := hubKey(sessionId, flowId)
	hubRegMu.Lock()
	defer hubRegMu.Unlock()
	if h, ok := hubReg[key]; ok {
		return h, false
	}
	hubReg[key] = hub
	return hub, true
}

// ====================== 订阅 / 发布 ======================

func (h *execHub) Subscribe(conn *wsCommon.WsConnection) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.subs[conn] = struct{}{}
}

func (h *execHub) Unsubscribe(conn *wsCommon.WsConnection) {
	h.mu.Lock()
	defer h.mu.Unlock()
	delete(h.subs, conn)
}

// Publish 广播消息到所有订阅连接（跳过已关闭连接；WriteJSON 自带写锁，并发安全）
func (h *execHub) Publish(msg *wsCommon.WsPushMsg) {
	h.mu.Lock()
	conns := make([]*wsCommon.WsConnection, 0, len(h.subs))
	for c := range h.subs {
		conns = append(conns, c)
	}
	h.mu.Unlock()
	for _, c := range conns {
		if c.IsClosed() {
			h.Unsubscribe(c)
			continue
		}
		_ = writeJSON(c, msg)
	}
}

// ReportStart / ReportComplete 实现 ProgressReporter：节点进度广播
func (h *execHub) ReportStart(nodeId, nodeName string, nodeIndex, nodeCount int) {
	h.Publish(&wsCommon.WsPushMsg{
		Type:    "node_start",
		Message: fmt.Sprintf("开始执行(%d/%d): %s ", nodeIndex, nodeCount, nodeName),
		Data: map[string]interface{}{
			"nodeId": nodeId, "nodeName": nodeName,
			"nodeIndex": nodeIndex, "nodeCount": nodeCount,
		},
	})
}

func (h *execHub) ReportComplete(nodeId, nodeName string, nodeIndex, nodeCount int) {
	h.Publish(&wsCommon.WsPushMsg{
		Type:    "node_complete",
		Message: fmt.Sprintf("执行完成(%d/%d): %s ", nodeIndex, nodeCount, nodeName),
		Data: map[string]interface{}{
			"nodeId": nodeId, "nodeName": nodeName,
			"nodeIndex": nodeIndex, "nodeCount": nodeCount,
		},
	})
}

// ====================== 接管 / 取消 ======================

// SetCancel 预留取消函数（候选 hub 在尚未确定执行方时设置，供用户提前取消）
func (h *execHub) SetCancel(cancel context.CancelFunc) {
	h.mu.Lock()
	h.cancel = cancel
	h.mu.Unlock()
}

// MarkOwned 标记执行方已接管本 hub（真正开始 BuildExecution 前调用）
func (h *execHub) MarkOwned() {
	h.mu.Lock()
	h.owned = true
	h.mu.Unlock()
}

// TryOwn 原子接管：仅当未被持有才置 owned 并设 cancel，返回是否抢到所有权。
// 并发恢复（扫描/用户附着）用它在锁/DB 前置位，避免"检查 Owned→MarkOwned"竞态下
// 后到者覆盖先到者的 cancel（用户取消会取消错 ctx）。失败者只附着订阅、不重复拉起执行。
func (h *execHub) TryOwn(cancel context.CancelFunc) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.owned {
		return false
	}
	h.owned = true
	h.cancel = cancel
	return true
}

func (h *execHub) Owned() bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.owned
}

// CancelByUser 用户显式取消（workflow_cancel）：置 userCancelled 标志并取消执行。
// 恢复例程据此把终态写成"用户已终止执行 / retryable=0"（永久取消，不再被扫描捞起）。
func (h *execHub) CancelByUser() {
	h.userCancelled.Store(true)
	h.mu.Lock()
	c := h.cancel
	h.mu.Unlock()
	if c != nil {
		c()
	}
}

func (h *execHub) UserCancelled() bool {
	return h.userCancelled.Load()
}

// Close 关闭 hub：注销注册、退订全部连接、清其 meta、close done（sync.Once，可被
// 执行方与占用方重复调用）。仅当连接 meta 仍指向本 hub 时清空（指针比较防误清新订阅）。
func (h *execHub) Close() {
	h.closeOnce.Do(func() {
		key := hubKey(h.sessionId, h.flowId)
		hubRegMu.Lock()
		if hubReg[key] == h {
			delete(hubReg, key)
		}
		hubRegMu.Unlock()

		h.mu.Lock()
		conns := make([]*wsCommon.WsConnection, 0, len(h.subs))
		for c := range h.subs {
			conns = append(conns, c)
		}
		h.closed = true
		h.mu.Unlock()

		for _, c := range conns {
			if cur, ok := wsCommon.GetMetaT[*execHub](c, "execHub"); ok && cur == h {
				c.SetMeta("execHub", nil)
				c.SetMeta("execCancel", nil)
			}
		}
		close(h.done)
	})
}

// Done 返回 hub 关闭通知 channel（供订阅 watcher 等退出）
func (h *execHub) Done() <-chan struct{} {
	return h.done
}

// getProgressHub 从 context 取 hub（reporter 即 hub 本身）
func getProgressHub(ctx context.Context) *execHub {
	if h, ok := GetProgressReporter(ctx).(*execHub); ok {
		return h
	}
	return nil
}

// subscribeConnToHub 把连接挂到 hub 上并接入取消：
// conn meta execCancel = hub.CancelByUser → 现有 workflow_cancel 处理器（handleCancel）直接生效。
// 断开连接不触发取消（用户执行经 closeCtx→execCtx 既有链取消；恢复执行脱离连接，订阅者仅观察）。
// 注意：必须转成 context.CancelFunc 再存，否则 GetMetaT[context.CancelFunc] 的类型断言
// （动态类型须与命名类型完全一致）会因方法值类型为 func() 而失败，取消静默失效。
func subscribeConnToHub(conn *wsCommon.WsConnection, hub *execHub) {
	hub.Subscribe(conn)
	conn.SetMeta("execHub", hub)
	conn.SetMeta("execCancel", context.CancelFunc(hub.CancelByUser))
}

// execAttach 恢复例程附着的用户连接（用户点击执行发现陈旧运行中记录时传入）。
// hub 为连接已订阅的候选 hub：恢复例程复用同一实例接管（registerHubIfAbsent 按 session+flow 去重，
// 全进程同一时刻至多一个 hub 实例），避免"占位被关→重建"竞态导致连接订阅/取消丢失。
type execAttach struct {
	conn      *wsCommon.WsConnection
	sessionId string
	flowId    int64
	hub       *execHub // 可能为 nil（调用方无候选时恢复例程自行注册）
}