cnb/agent/routes/workspace/keep.ts

import { createSkill, tool } from '@kevisual/router';
import { app, cnb } from '../../app.ts';
import { nanoid } from 'nanoid';
import dayjs from 'dayjs';
import { createKeepAlive } from '../../../src/keep.ts';

type AliveInfo = {
  startTime: number;
  updatedTime?: number;
  KeepAlive: ReturnType<typeof createKeepAlive>;
  id: string;// 6位唯一标识符
}

const keepAliveMap = new Map<string, AliveInfo>();

// 保持工作空间存活技能
app.route({
  path: 'cnb',
  key: 'keep-workspace-alive',
  description: '保持工作空间存活技能,参数wsUrl:工作空间访问URL，cookie:访问工作空间所需的cookie',
  middleware: ['admin-auth'],
  metadata: {
    tags: [],
    ...({
      args: {
        wsUrl: tool.schema.string().describe('工作空间的访问URL'),
        cookie: tool.schema.string().describe('访问工作空间所需的cookie')
      }
    })
  }
}).define(async (ctx) => {
  const wsUrl = ctx.query?.wsUrl as string;
  const cookie = ctx.query?.cookie as string;
  if (!wsUrl) {
    ctx.throw(400, '缺少工作空间访问URL参数');
  }
  if (!cookie) {
    ctx.throw(400, '缺少访问工作空间所需的cookie参数');
  }

  // 检测是否已在运行（通过 wsUrl 遍历检查）
  const existing = Array.from(keepAliveMap.values()).find(info => (info as AliveInfo).id && (info as any).KeepAlive?.wsUrl === wsUrl);
  if (existing) {
    ctx.body = { message: `工作空间 ${wsUrl} 的保持存活任务已在运行中`, id: (existing as AliveInfo).id };
    return;
  }

  console.log(`启动保持工作空间 ${wsUrl} 存活的任务`);
  const keep = createKeepAlive({
    wsUrl,
    cookie,
    onConnect: () => {
      console.log(`工作空间 ${wsUrl} 保持存活任务已连接`);
    },
    onMessage: (data) => {
      // 可选：处理收到的消息
      // console.log(`工作空间 ${wsUrl} 收到消息: ${data}`);
      // 通过 wsUrl 找到对应的 id 并更新时间
      for (const info of keepAliveMap.values()) {
        if ((info as any).KeepAlive?.wsUrl === wsUrl) {
          info.updatedTime = Date.now();
          break;
        }
      }
    },
    debug: true,
    onExit: (code) => {
      console.log(`工作空间 ${wsUrl} 保持存活任务已退出，退出码: ${code}`);
      // 通过 wsUrl 找到对应的 id 并删除
      for (const [id, info] of keepAliveMap.entries()) {
        if ((info as any).KeepAlive?.wsUrl === wsUrl) {
          keepAliveMap.delete(id);
          break;
        }
      }
    }
  });

  const id = nanoid(6).toLowerCase();
  keepAliveMap.set(id, { startTime: Date.now(), updatedTime: Date.now(), KeepAlive: keep, id });

  ctx.body = { content: `已启动保持工作空间 ${wsUrl} 存活的任务`, id };
}).addTo(app);

// 获取保持工作空间存活任务列表技能
app.route({
  path: 'cnb',
  key: 'list-keep-alive-tasks',
  description: '获取保持工作空间存活任务列表技能',
  middleware: ['admin-auth'],
  metadata: {
    tags: [],
  }
}).define(async (ctx) => {
  const list = Array.from(keepAliveMap.entries()).map(([id, info]) => {
    const now = Date.now();
    const duration = Math.floor((now - info.startTime) / 60000); // 分钟
    return {
      id,
      wsUrl: (info as any).KeepAlive?.wsUrl,
      startTime: info.startTime,
      startTimeStr: dayjs(info.startTime).format('YYYY-MM-DD HH:mm'),
      updatedTime: info.updatedTime,
      updatedTimeStr: dayjs(info.updatedTime).format('YYYY-MM-DD HH:mm'),
      duration,
    }
  });
  ctx.body = { list };
}).addTo(app);

// 停止保持工作空间存活技能
app.route({
  path: 'cnb',
  key: 'stop-keep-workspace-alive',
  description: '停止保持工作空间存活技能, 参数wsUrl:工作空间访问URL或者id',
  middleware: ['admin-auth'],
  metadata: {
    tags: [],
    ...({
      args: {
        wsUrl: tool.schema.string().optional().describe('工作空间的访问URL'),
        id: tool.schema.string().optional().describe('保持存活任务的唯一标识符'),
      }
    })
  }
}).define(async (ctx) => {
  const wsUrl = ctx.query?.wsUrl as string;
  const id = ctx.query?.id as string;
  if (!wsUrl && !id) {
    ctx.throw(400, '缺少工作空间访问URL参数或唯一标识符');
  }

  let targetId: string | undefined;
  let wsUrlFound: string | undefined;

  if (id) {
    const info = keepAliveMap.get(id);
    if (info) {
      targetId = id;
      wsUrlFound = (info as any).KeepAlive?.wsUrl;
    }
  } else if (wsUrl) {
    for (const [key, info] of keepAliveMap.entries()) {
      if ((info as any).KeepAlive?.wsUrl === wsUrl) {
        targetId = key;
        wsUrlFound = wsUrl;
        break;
      }
    }
  }

  if (targetId) {
    const keepAlive = keepAliveMap.get(targetId);
    const endTime = Date.now();
    const duration = endTime - keepAlive!.startTime;
    keepAlive?.KeepAlive?.disconnect();
    keepAliveMap.delete(targetId);
    ctx.body = { content: `已停止保持工作空间 ${wsUrlFound} 存活的任务，持续时间: ${duration}ms`, id: targetId };
  } else {
    ctx.body = { content: `没有找到对应的工作空间保持存活任务` };
  }
}).addTo(app);

app.route({
  path: 'cnb',
  key: 'reset-keep-workspace-alive',
  description: '对存活的工作空间，startTime进行重置',
  middleware: ['admin-auth'],
  metadata: {
    tags: [],
  }
}).define(async (ctx) => {
  const now = Date.now();
  for (const info of keepAliveMap.values()) {
    info.startTime = now;
  }
  ctx.body = { content: `已重置所有存活工作空间的开始时间` };
}).addTo(app);

app.route({
  path: 'cnb',
  key: 'clear-keep-workspace-alive',
  description: '对存活的工作空间，超过5小时的进行清理',
  middleware: ['admin-auth'],
  metadata: {
    tags: [],
  }
}).define(async (ctx) => {
  const res = clearKeepAlive();
  ctx.body = {
    content: `已清理所有存活工作空间中超过5小时的任务` + (res.length ? `，清理项：${res.map(i => i.wsUrl).join(', ')}` : ''),
    list: res
  };
}).addTo(app);

const clearKeepAlive = () => {
  const now = Date.now();
  let clearedArr: { id: string; wsUrl: string }[] = [];
  for (const [id, info] of keepAliveMap.entries()) {
    if (now - info.startTime > FIVE_HOURS) {
      console.log(`工作空间 ${(info as any).KeepAlive?.wsUrl} 超过5小时，自动停止`);
      info.KeepAlive?.disconnect?.();
      keepAliveMap.delete(id);
      clearedArr.push({ id, wsUrl: (info as any).KeepAlive?.wsUrl });
    }
  }
  return clearedArr;
}

// 每5小时自动清理超时的keepAlive任务
const FIVE_HOURS = 5 * 60 * 60 * 1000;
setInterval(() => {
  clearKeepAlive();
}, FIVE_HOURS);