shuzhiren-comfyui/message-dispatcher/src/api/index.js

import express from 'express';
import bridgeManager from '../bridge-manager/index.js';
import websocketServer from '../websocket-server/index.js';
import taskScheduler from '../task-scheduler/index.js';
import { v4 as uuidv4 } from 'uuid';
import logger from '../logger/index.js';
import { authMiddleware } from '../auth/middleware.js';

const router = express.Router();

router.get('/capacity', authMiddleware, (req, res) => {
  const capacity = bridgeManager.getAvailableCapacity();
  const stats = taskScheduler.getStats();
  
  res.json({
    success: true,
    data: {
      ...capacity,
      processingTasks: stats.processing,
      pendingTasks: stats.pending,
      availableForNewTasks: Math.max(0, capacity.available - stats.pending)
    }
  });
});

router.get('/health', (req, res) => {
  res.json({
    success: true,
    data: {
      status: 'ok',
      timestamp: new Date().toISOString()
    }
  });
});

router.get('/bridges', authMiddleware, (req, res) => {
  const bridges = bridgeManager.getAllBridges();
  res.json({
    success: true,
    data: bridges
  });
});

router.get('/bridges/:bridgeId', authMiddleware, (req, res) => {
  const bridge = bridgeManager.getBridge(req.params.bridgeId);
  if (!bridge) {
    return res.status(404).json({
      success: false,
      error: '桥接器不存在'
    });
  }
  res.json({
    success: true,
    data: {
      id: bridge.id,
      info: bridge.info,
      connectedAt: bridge.connectedAt,
      lastHeartbeat: bridge.lastHeartbeat
    }
  });
});

router.post('/task', authMiddleware, async (req, res) => {
  try {
    const { workflowId, nodeInfoList, webhookUrl, taskType } = req.body;
    const requestId = uuidv4();

    const capacity = bridgeManager.getAvailableCapacity();
    await taskScheduler.setCurrentCapacity(capacity.online, false);

    const assignResult = taskScheduler.tryDirectAssign({
      requestId,
      workflowId,
      nodeInfoList,
      webhookUrl,
      taskType
    });

    if (assignResult.success) {
      logger.info(`收到任务请求, 直接分配实例: ${assignResult.instanceId}, bridgeId: ${assignResult.bridgeId}, requestId: ${requestId}`);
      
      try {
        const result = await websocketServer.sendTaskToInstance(
          assignResult.bridgeId,
          assignResult.instanceId,
          { workflowId, nodeInfoList, webhookUrl },
          requestId
        );

        res.json({
          success: true,
          data: {
            requestId,
            instanceId: assignResult.instanceId,
            bridgeId: assignResult.bridgeId,
            status: 'processing',
            ...result.data
          }
        });
      } catch (sendError) {
        logger.error('发送任务失败:', sendError);
        taskScheduler.handleTaskFailure(requestId, sendError.message);
        
        res.status(500).json({
          success: false,
          error: sendError.message,
          requestId
        });
      }
    } else {
      const queueResult = taskScheduler.addTask({
        requestId,
        workflowId,
        nodeInfoList,
        webhookUrl,
        taskType
      });

      if (queueResult.success) {
        logger.info(`收到任务请求, 无可用实例, 加入队列: ${requestId}, 队列位置: ${queueResult.queuePosition}`);
        
        res.json({
          success: true,
          data: {
            requestId,
            status: 'queued',
            queuePosition: queueResult.queuePosition,
            capacity: capacity
          }
        });
      } else {
        logger.warn(`任务队列已满: ${requestId}`);
        
        res.status(503).json({
          success: false,
          error: queueResult.error,
          data: {
            requestId,
            capacity
          }
        });
      }
    }
  } catch (error) {
    logger.error('处理任务请求失败:', error);
    res.status(500).json({
      success: false,
      error: error.message
    });
  }
});

router.get('/task/:taskId', authMiddleware, (req, res) => {
  const taskStatus = taskScheduler.getTaskStatus(req.params.taskId);
  
  if (!taskStatus) {
    return res.status(404).json({
      success: false,
      error: '任务不存在'
    });
  }
  
  res.json({
    success: true,
    data: taskStatus
  });
});

router.get('/instances', authMiddleware, (req, res) => {
  const instances = bridgeManager.getAllInstances();

  res.json({
    success: true,
    data: instances
  });
});

router.get('/overview', authMiddleware, (req, res) => {
  const bridges = bridgeManager.getAllBridges();
  const capacity = bridgeManager.getAvailableCapacity();
  const stats = taskScheduler.getStats();
  
  let totalInstances = 0;
  let onlineInstances = 0;
  let busyInstances = 0;
  let offlineInstances = 0;

  for (const bridge of bridges) {
    if (bridge.info?.instances) {
      totalInstances += bridge.info.instances.length;
      onlineInstances += bridge.info.instances.filter(i => i.status === 'online').length;
      busyInstances += bridge.info.instances.filter(i => i.status === 'busy').length;
      offlineInstances += bridge.info.instances.filter(i => i.status === 'offline').length;
    }
  }

  res.json({
    success: true,
    data: {
      bridges: {
        total: bridges.length,
        online: bridges.length
      },
      instances: {
        total: totalInstances,
        online: onlineInstances,
        busy: busyInstances,
        offline: offlineInstances,
        locked: capacity.locked,
        available: capacity.available
      },
      tasks: stats
    }
  });
});

router.get('/monitor/overview', authMiddleware, (req, res) => {
  const bridges = bridgeManager.getAllBridges();
  const capacity = bridgeManager.getAvailableCapacity();
  const stats = taskScheduler.getStats();
  
  let totalInstances = 0;
  let onlineInstances = 0;
  let busyInstances = 0;
  let offlineInstances = 0;

  for (const bridge of bridges) {
    if (bridge.info?.instances) {
      totalInstances += bridge.info.instances.length;
      onlineInstances += bridge.info.instances.filter(i => i.status === 'online').length;
      busyInstances += bridge.info.instances.filter(i => i.status === 'busy').length;
      offlineInstances += bridge.info.instances.filter(i => i.status === 'offline').length;
    }
  }

  res.json({
    success: true,
    data: {
      instances: {
        total: totalInstances,
        online: onlineInstances,
        busy: busyInstances,
        offline: offlineInstances,
        locked: capacity.locked,
        available: capacity.available
      },
      tasks: stats
    }
  });
});

router.get('/tasks', authMiddleware, (req, res) => {
  const stats = taskScheduler.getStats();
  
  res.json({
    success: true,
    data: {
      stats,
      pending: taskScheduler.pendingTaskQueue.slice(0, 20),
      processing: Array.from(taskScheduler.processingTasks.values()),
      recovering: Array.from(taskScheduler.recoveringTasks.values())
    }
  });
});

router.post('/instances/:instanceId/health-check', authMiddleware, async (req, res) => {
  try {
    const { instanceId } = req.params;
    const bridges = bridgeManager.getAllBridges();
    
    let targetBridgeId = null;
    for (const bridge of bridges) {
      if (bridge.info?.instances) {
        const instance = bridge.info.instances.find(i => i.id === instanceId);
        if (instance) {
          targetBridgeId = bridge.id;
          break;
        }
      }
    }

    if (!targetBridgeId) {
      return res.status(404).json({
        success: false,
        error: '实例不存在'
      });
    }

    const requestId = uuidv4();
    logger.info(`收到实例健康检查请求, instanceId: ${instanceId}, bridgeId: ${targetBridgeId}`);

    const result = await websocketServer.sendInstanceCheckToBridge(
      targetBridgeId,
      'single',
      instanceId,
      requestId
    );

    res.json({
      success: true,
      data: result.data
    });
  } catch (error) {
    logger.error('处理实例健康检查失败:', error);
    res.status(500).json({
      success: false,
      error: error.message
    });
  }
});

router.post('/instances/health-check-offline', authMiddleware, async (req, res) => {
  try {
    const bridges = bridgeManager.getAllBridges();
    const results = [];

    for (const bridge of bridges) {
      const requestId = uuidv4();
      logger.info(`收到离线实例检查请求, bridgeId: ${bridge.id}`);

      try {
        const result = await websocketServer.sendInstanceCheckToBridge(
          bridge.id,
          'offline',
          null,
          requestId
        );
        results.push({
          bridgeId: bridge.id,
          success: true,
          data: result.data
        });
      } catch (error) {
        results.push({
          bridgeId: bridge.id,
          success: false,
          error: error.message
        });
      }
    }

    res.json({
      success: true,
      data: results
    });
  } catch (error) {
    logger.error('处理离线实例检查失败:', error);
    res.status(500).json({
      success: false,
      error: error.message
    });
  }
});

router.post('/instances/health-check-all', authMiddleware, async (req, res) => {
  try {
    const bridges = bridgeManager.getAllBridges();
    const results = [];

    for (const bridge of bridges) {
      const requestId = uuidv4();
      logger.info(`收到全部实例检查请求, bridgeId: ${bridge.id}`);

      try {
        const result = await websocketServer.sendInstanceCheckToBridge(
          bridge.id,
          'all',
          null,
          requestId
        );
        results.push({
          bridgeId: bridge.id,
          success: true,
          data: result.data
        });
      } catch (error) {
        results.push({
          bridgeId: bridge.id,
          success: false,
          error: error.message
        });
      }
    }

    res.json({
      success: true,
      data: results
    });
  } catch (error) {
    logger.error('处理全部实例检查失败:', error);
    res.status(500).json({
      success: false,
      error: error.message
    });
  }
});

export default router;