shalu-agent-workflow/apps/web/src/store/modules/runner.store.ts

import { defineStore } from 'pinia'
import { computed, reactive, readonly, ref } from 'vue'
import { dayjs } from 'element-plus'

export type RunnerStatus = 'idle' | 'connecting' | 'running' | 'finished' | 'error' | 'suspended'

export type NodeStatus = 'idle' | 'running' | 'success' | 'failed' | 'suspended'

export interface RunnerNodeInfo {
	nodeId: string
	nodeName: string
	nodeType: string
}

export interface RunnerNodeState extends RunnerNodeInfo {
	status: NodeStatus
	lastUpdateTime?: string
	track?: any
}

export interface RunnerExecution {
	runnerKey: string
	status: RunnerStatus
	startedAt: string | null
	finishedAt: string | null
	nodes: RunnerNodeState[]
	result?: any
}

interface AgentRunnerMessageBase {
	cmd: string
	time?: string
}

interface ConnectErrorMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_CONNECT_ERROR_MSG'
	errorMsg: string
}

interface WelcomeMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_WELCOME_MSG'
}

interface HeartbeatMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_HEARTBEAT_MSG'
}

interface AgentRunningMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_AGENT_RUNNING_MSG'
}

interface NodeRunningMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_NODE_RUNNING_MSG'
	node: RunnerNodeInfo
}

interface NodeFinishMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_NODE_FINISH_MSG'
	node: RunnerNodeInfo
	track: any
}

interface NodeIterationRunningMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_NODE_ITERATION_RUNNING_MSG'
	count: number
	node: RunnerNodeInfo
}

interface NodeIterationStepMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_NODE_ITERATION_STEP_MSG'
	count: number
	index: number
	node: RunnerNodeInfo
}

interface NodeIterationFinishMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_NODE_ITERATION_FINISH_MSG'
	node: RunnerNodeInfo
}

interface AgentFinishMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_AGENT_FINISH_MSG'
	result: any
}

interface AgentErrorMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_AGENT_ERROR_MSG'
	errorMsg: string
}

interface AgentSuspendMessage extends AgentRunnerMessageBase {
	cmd: 'CMD_AGENT_SUSPEND_MSG'
}

export type AgentRunnerMessage =
	| ConnectErrorMessage
	| WelcomeMessage
	| HeartbeatMessage
	| AgentRunningMessage
	| NodeRunningMessage
	| NodeFinishMessage
	| NodeIterationRunningMessage
	| NodeIterationStepMessage
	| NodeIterationFinishMessage
	| AgentFinishMessage
	| AgentErrorMessage
	| AgentSuspendMessage

const AGENT_RUNNER_WS_BASE = `wss://${import.meta.env.VITE_BASE_URL}/api/ws/agentRunner`

export const useRunnerStore = defineStore('runner', () => {
	const currentRunnerKey = ref<string | null>(null)
	const currentStartNodeId = ref<string | null>(null)
	const status = ref<RunnerStatus>('idle')
	const errorMsg = ref<string | null>(null)
	const connected = ref(false)
	const lastHeartbeatAt = ref<number | null>(null)
	const agentResult = ref<any>(null)

	const nodesMap = reactive<Record<string, RunnerNodeState>>({})

	const executions = ref<RunnerExecution[]>([])

	const socket = ref<WebSocket | null>(null)
	const heartbeatTimer = ref<number | null>(null)

	const nodes = computed(() => Object.values(nodesMap))

	const getCurrentExecution = () => {
		if (!currentRunnerKey.value) return null
		return executions.value.find((item) => item.runnerKey === currentRunnerKey.value) || null
	}

	const clearHeartbeat = () => {
		if (heartbeatTimer.value) {
			window.clearInterval(heartbeatTimer.value)
			heartbeatTimer.value = null
		}
	}

	const resetState = () => {
		status.value = 'idle'
		errorMsg.value = null
		connected.value = false
		lastHeartbeatAt.value = null
		agentResult.value = null
		Object.keys(nodesMap).forEach((key) => {
			delete nodesMap[key]
		})
	}

	const closeSocket = () => {
		clearHeartbeat()
		if (socket.value) {
			try {
				socket.value.close()
			} catch {
				// ignore
			}
			socket.value = null
		}
	}

	const updateNodeState = (info: RunnerNodeInfo, partial: Partial<RunnerNodeState>) => {
		const existing = nodesMap[info.nodeId]
		const nextState: RunnerNodeState = {
			nodeId: info.nodeId,
			nodeName: info.nodeName,
			nodeType: info.nodeType,
			status: existing?.status ?? 'idle',
			...existing,
			...partial
		}

		nodesMap[info.nodeId] = nextState

		const execution = getCurrentExecution()
		if (execution) {
			const index = execution.nodes.findIndex((item) => item.nodeId === info.nodeId)
			if (index === -1) {
				execution.nodes.push({ ...nextState })
			} else {
				execution.nodes[index] = { ...nextState }
			}
		}
	}

	const handleMessage = (raw: MessageEvent<string>) => {
		let data: AgentRunnerMessage | null = null
		try {
			data = JSON.parse(raw.data)
		} catch {
			return
		}

		if (!data || typeof data !== 'object' || !('cmd' in data)) return

		switch (data.cmd) {
			/**
			 * 连接运行器失败
			 */
			case 'CMD_CONNECT_ERROR_MSG': {
				const msg = data as ConnectErrorMessage
				status.value = 'error'
				errorMsg.value = msg.errorMsg || '连接运行器失败'
				connected.value = false

				const execution = getCurrentExecution()
				if (execution) {
					execution.status = 'error'
					execution.finishedAt = msg.time || new Date().toISOString()
				}

				closeSocket()
				break
			}
			/**
			 * 连接运行器成功
			 */
			case 'CMD_WELCOME_MSG': {
				status.value = 'running'
				connected.value = true

				const execution = getCurrentExecution()
				if (execution) {
					execution.status = 'running'
				}

				break
			}
			/**
			 * 心跳消息
			 */
			case 'CMD_HEARTBEAT_MSG': {
				lastHeartbeatAt.value = Date.now()
				break
			}
			/**
			 * 智能体运行消息
			 */
			case 'CMD_AGENT_RUNNING_MSG': {
				status.value = 'running'

				const execution = getCurrentExecution()
				if (execution) {
					execution.status = 'running'
				}

				break
			}
			/**
			 * 节点运行消息
			 */
			case 'CMD_NODE_RUNNING_MSG': {
				const msg = data as NodeRunningMessage
				updateNodeState(msg.node, {
					status: 'running',
					lastUpdateTime: msg.time
				})
				break
			}
			/**
			 * 节点运行完成消息
			 */
			case 'CMD_NODE_FINISH_MSG': {
				const msg = data as NodeFinishMessage
				const isSuccess = !!msg.track?.is_success
				updateNodeState(msg.node, {
					status: isSuccess ? 'success' : 'failed',
					lastUpdateTime: msg.time,
					track: msg.track
				})
				break
			}
			/**
			 * 节点迭代运行消息
			 */
			case 'CMD_NODE_ITERATION_RUNNING_MSG': {
				const msg = data as NodeIterationRunningMessage
				updateNodeState(msg.node, {
					status: 'running',
					lastUpdateTime: msg.time
				})
				break
			}
			/**
			 * 节点迭代运行步骤消息
			 */
			case 'CMD_NODE_ITERATION_STEP_MSG': {
				const msg = data as NodeIterationStepMessage
				updateNodeState(msg.node, {
					status: 'running',
					lastUpdateTime: msg.time
				})
				break
			}

			/**
			 * 节点迭代运行完成消息
			 */
			case 'CMD_NODE_ITERATION_FINISH_MSG': {
				const msg = data as NodeIterationFinishMessage
				updateNodeState(msg.node, {
					status: 'success',
					lastUpdateTime: msg.time
				})
				break
			}

			/**
			 * 智能体运行完成消息
			 */
			case 'CMD_AGENT_FINISH_MSG': {
				const msg = data as AgentFinishMessage
				status.value = 'finished'
				agentResult.value = msg.result

				const execution = getCurrentExecution()
				if (execution) {
					execution.status = 'finished'
					execution.finishedAt = msg.time || new Date().toISOString()
					execution.result = msg.result
				}

				closeSocket()
				break
			}

			/**
			 * 智能体运行异常失败
			 */
			case 'CMD_AGENT_ERROR_MSG': {
				const msg = data as AgentErrorMessage
				status.value = 'error'
				agentResult.value = msg.errorMsg || '智能体运行异常失败'

				const execution = getCurrentExecution()
				if (execution) {
					execution.status = 'finished'
					execution.finishedAt = msg.time || new Date().toISOString()
					execution.result = msg.errorMsg || '智能体运行异常失败'
				}

				closeSocket()
				break
			}

			/**
			 * 智能体运行被挂起
			 */
			case 'CMD_AGENT_SUSPEND_MSG': {
				status.value = 'suspended'

				const execution = getCurrentExecution()
				if (execution) {
					execution.status = 'suspended'
				}

				Object.values(nodesMap).forEach((node) => {
					if (node.status === 'running') {
						updateNodeState(node, {
							status: 'suspended',
							lastUpdateTime: data?.time || node.lastUpdateTime
						})
					}
				})
				break
			}
			default:
				break
		}
	}

	const startHeartbeat = () => {
		clearHeartbeat()
		heartbeatTimer.value = window.setInterval(() => {
			if (!socket.value || socket.value.readyState !== WebSocket.OPEN) return
			socket.value.send(JSON.stringify({ cmd: 'CMD_HEARTBEAT_MSG' }))
		}, 10000)
	}

	const startRunner = (runnerKey: string, startNodeId?: string) => {
		if (!runnerKey) return

		closeSocket()
		resetState()

		currentRunnerKey.value = runnerKey
		currentStartNodeId.value = startNodeId || null
		status.value = 'connecting'

		executions.value.unshift({
			runnerKey,
			status: 'connecting',
			startedAt: dayjs().format('YYYY-MM-DD HH:mm:ss'),
			finishedAt: null,
			nodes: [],
			result: undefined
		})

		const url = `${AGENT_RUNNER_WS_BASE}?agentRunnerKey=${encodeURIComponent(runnerKey)}`
		const ws = new WebSocket(url)
		socket.value = ws

		ws.onopen = () => {
			connected.value = true
			status.value = 'running'
			startHeartbeat()

			const execution = getCurrentExecution()
			if (execution) {
				execution.status = 'running'
			}
		}

		ws.onmessage = (event) => {
			handleMessage(event as MessageEvent<string>)
		}

		ws.onerror = () => {
			status.value = 'error'
			errorMsg.value = errorMsg.value || '运行器连接异常'

			const execution = getCurrentExecution()
			if (execution) {
				execution.status = 'error'
				if (!execution.finishedAt) {
					execution.finishedAt = dayjs().format('YYYY-MM-DD HH:mm:ss')
				}
			}
		}

		ws.onclose = () => {
			connected.value = false
			clearHeartbeat()
			if (status.value === 'running' || status.value === 'connecting') {
				status.value = 'finished'

				const execution = getCurrentExecution()
				if (execution && !execution.finishedAt) {
					execution.status = 'finished'
					execution.finishedAt = dayjs().format('YYYY-MM-DD HH:mm:ss')
				}
			}
		}
	}

	const stopRunner = () => {
		closeSocket()
		status.value = 'finished'

		const execution = getCurrentExecution()
		if (execution && !execution.finishedAt) {
			execution.status = 'finished'
			execution.finishedAt = dayjs().format('YYYY-MM-DD HH:mm:ss')
		}
	}

	const resetRunner = () => {
		closeSocket()
		currentRunnerKey.value = null
		currentStartNodeId.value = null
		resetState()
	}

	return {
		currentRunnerKey: readonly(currentRunnerKey),
		currentStartNodeId: readonly(currentStartNodeId),
		status,
		errorMsg,
		connected,
		lastHeartbeatAt,
		nodes,
		agentResult,
		executions,
		startRunner,
		stopRunner,
		resetRunner
	}
})