aiquery2/api/app/controllers/superAgentProtocol.js

/**
 * 文件作用：
 * 为 SuperAgent v2 useStream 协议提供经过认证的透明代理。
 *
 * 职责：
 * 1. 从 tenderUserAuth 写入的上下文解析可信用户身份。
 * 2. 转发固定的 thread state、history、commands、cancel JSON 请求。
 * 3. 以无缓冲方式转发 stream/events SSE，并在浏览器断开时释放上游连接。
 *
 * 不负责：
 * 1. 兼容旧 /api/chat、checkpoint 或 runs/stream 协议。
 * 2. 创建任务、解释 Agent 协议内容或修改上游响应 envelope。
 *
 * 维护说明：
 * 只能通过本文件导出的固定 handler 构造上游路径，禁止接受客户端提供的目标地址或转发 Header。
 */
'use strict';

const { Readable } = require('node:stream');

const CONNECT_TIMEOUT_MS = 10_000;
const JSON_TIMEOUT_MS = 30_000;
const ALLOWED_CANCEL_QUERY = ['wait', 'action'];

const resolveAuthenticatedUserId = (ctx) => {
   const userInfo = ctx.fs?.curUser?.userInfo;
   if (!userInfo || typeof userInfo !== 'object') return '';
   const value =
      userInfo.localUserId ??
      userInfo.id ??
      userInfo.userId ??
      userInfo.pepUserId ??
      userInfo.pep_user_id ??
      userInfo.pepId ??
      userInfo.pep_id ??
      '';
   return String(value).trim();
};

const getUpstreamBaseUrl = (ctx) => {
   const baseUrl = String(ctx.app.fs.config?.superAgent?.baseUrl || '').trim();
   if (!baseUrl) throw new Error('未配置 superAgent.baseUrl');
   return baseUrl.replace(/\/+$/, '');
};

const buildTrustedHeaders = (ctx, hasJsonBody = false) => {
   const userId = resolveAuthenticatedUserId(ctx);
   if (!userId) return null;
   const headers = {
      Accept: 'application/json, text/event-stream',
      'X-User-Id': userId,
   };
   const searchToken = String(ctx.fs?.curUser?.token || '').trim();
   if (searchToken) headers['X-Search-Token'] = searchToken;
   if (hasJsonBody) headers['Content-Type'] = 'application/json';
   return headers;
};

const buildThreadPath = (threadId, suffix = '') => {
   return `/api/v1/threads/${encodeURIComponent(String(threadId || ''))}${suffix}`;
};

const buildCancelQuery = (ctx) => {
   const search = new URLSearchParams();
   ALLOWED_CANCEL_QUERY.forEach((key) => {
      const value = ctx.query?.[key];
      if (value !== undefined && value !== null && value !== '') {
         search.set(key, String(value));
      }
   });
   const query = search.toString();
   return query ? `?${query}` : '';
};

const writeSafeUpstreamError = async (ctx, response) => {
   const contentType = String(response.headers.get('content-type') || '');
   if (contentType.includes('application/json')) {
      try {
         ctx.body = await response.json();
         return;
      } catch (_) {
         // 统一进入安全错误，避免将无效上游内容或 HTML 返回给浏览器。
      }
   }
   ctx.body = { message: '智能体服务暂不可用' };
};

const logRequest = (ctx, level, message, fields) => {
   const logger = ctx.logger?.[level];
   if (typeof logger !== 'function') return;
   const safeThreadId = String(fields.threadId || '').replace(/[\r\n]/g, '').slice(0, 128);
   logger.call(ctx.logger, `[superAgentProtocol] ${message}，method：${fields.method}，path：${fields.path}，threadId：${safeThreadId}，status：${fields.status}，耗时：${fields.duration}ms`);
};

/**
 * 功能：透明转发 v2 JSON 请求。
 * 使用场景：state、history、commands 和 cancel 固定协议端点。
 *
 * 入参：Koa ctx、固定 method/path 构造结果。
 * 返回：上游状态码与 JSON envelope；204 不返回 body。
 *
 * 注意：连接等待最多10秒，请求总时长最多30秒，日志不记录请求正文和认证令牌。
 */
const proxyJson = async (ctx, method, path) => {
   const startedAt = Date.now();
   const threadId = String(ctx.params.threadId || '');
   const headers = buildTrustedHeaders(ctx, method !== 'GET');
   if (!headers) {
      ctx.status = 401;
      ctx.body = { message: '认证用户信息无效' };
      return;
   }

   const controller = new AbortController();
   const connectTimer = setTimeout(() => controller.abort(new Error('upstream connection timeout')), CONNECT_TIMEOUT_MS);
   const totalTimer = setTimeout(() => controller.abort(new Error('upstream request timeout')), JSON_TIMEOUT_MS);
   try {
      const response = await fetch(`${getUpstreamBaseUrl(ctx)}${path}`, {
         method,
         headers,
         body: method === 'GET' ? undefined : JSON.stringify(ctx.request.body || {}),
         signal: controller.signal,
      });
      clearTimeout(connectTimer);
      ctx.status = response.status;
      if (response.status === 204) {
         ctx.body = null;
      } else if (response.ok) {
         const contentType = String(response.headers.get('content-type') || '');
         if (!contentType.includes('application/json')) {
            ctx.status = 502;
            ctx.body = { message: '智能体协议响应格式错误' };
         } else {
            ctx.body = await response.json();
         }
      } else {
         await writeSafeUpstreamError(ctx, response);
      }
      logRequest(ctx, 'info', '上游请求完成', { method, path, threadId, status: ctx.status, duration: Date.now() - startedAt });
   } catch (error) {
      const isTimeout = controller.signal.aborted;
      ctx.status = isTimeout ? 504 : 502;
      ctx.body = { message: isTimeout ? '智能体服务请求超时' : '智能体服务暂不可用' };
      logRequest(ctx, 'error', '上游请求失败', { method, path, threadId, status: ctx.status, duration: Date.now() - startedAt });
   } finally {
      clearTimeout(connectTimer);
      clearTimeout(totalTimer);
   }
};

/**
 * 功能：无缓冲透明转发 v2 SSE 事件流。
 * 使用场景：useStream 订阅 thread stream/events。
 *
 * 入参：Koa ctx。
 * 返回：直接写入浏览器原始响应流。
 *
 * 注意：浏览器断开只中止当前订阅，不调用 Agent cancel。
 */
const streamEvents = async (ctx) => {
   const startedAt = Date.now();
   const threadId = String(ctx.params.threadId || '');
   const path = buildThreadPath(threadId, '/stream/events');
   const headers = buildTrustedHeaders(ctx, true);
   if (!headers) {
      ctx.status = 401;
      ctx.body = { message: '认证用户信息无效' };
      return;
   }

   const controller = new AbortController();
   const connectTimer = setTimeout(() => controller.abort(new Error('upstream connection timeout')), CONNECT_TIMEOUT_MS);
   let upstreamStream = null;
   let finished = false;
   const abortSubscription = () => {
      if (!finished) controller.abort(new Error('browser disconnected'));
   };

   try {
      const response = await fetch(`${getUpstreamBaseUrl(ctx)}${path}`, {
         method: 'POST',
         headers,
         body: JSON.stringify(ctx.request.body || {}),
         signal: controller.signal,
      });
      clearTimeout(connectTimer);
      if (!response.ok || !response.body) {
         ctx.status = response.status || 502;
         await writeSafeUpstreamError(ctx, response);
         return;
      }

      ctx.respond = false;
      ctx.res.writeHead(response.status, {
         'Content-Type': 'text/event-stream; charset=utf-8',
         'Cache-Control': 'no-cache, no-transform',
         Connection: 'keep-alive',
         'X-Accel-Buffering': 'no',
      });
      upstreamStream = Readable.fromWeb(response.body);
      ctx.req.once('aborted', abortSubscription);
      ctx.res.once('close', abortSubscription);

      await new Promise((resolve, reject) => {
         upstreamStream.once('end', resolve);
         upstreamStream.once('error', reject);
         upstreamStream.pipe(ctx.res);
      });
      finished = true;
      if (!ctx.res.writableEnded) ctx.res.end();
      logRequest(ctx, 'info', 'SSE订阅结束', { method: 'POST', path, threadId, status: response.status, duration: Date.now() - startedAt });
   } catch (error) {
      if (ctx.respond !== false && !ctx.res.headersSent) {
         ctx.status = controller.signal.aborted ? 504 : 502;
         ctx.body = { message: controller.signal.aborted ? '智能体服务连接超时' : '智能体服务暂不可用' };
      } else if (!ctx.res.writableEnded) {
         ctx.res.end();
      }
      logRequest(ctx, 'error', 'SSE订阅异常', { method: 'POST', path, threadId, status: ctx.status || 502, duration: Date.now() - startedAt });
   } finally {
      finished = true;
      clearTimeout(connectTimer);
      ctx.req.removeListener('aborted', abortSubscription);
      ctx.res.removeListener('close', abortSubscription);
      if (upstreamStream && !upstreamStream.destroyed) upstreamStream.destroy();
   }
};

module.exports = {
   resolveAuthenticatedUserId,
   getState: (ctx) => proxyJson(ctx, 'GET', buildThreadPath(ctx.params.threadId, '/state')),
   getHistory: (ctx) => proxyJson(ctx, 'POST', buildThreadPath(ctx.params.threadId, '/history')),
   sendCommand: (ctx) => proxyJson(ctx, 'POST', buildThreadPath(ctx.params.threadId, '/commands')),
   streamEvents,
   cancelRun: (ctx) => proxyJson(
      ctx,
      'POST',
      `${buildThreadPath(ctx.params.threadId, `/runs/${encodeURIComponent(String(ctx.params.runId || ''))}/cancel`)}${buildCancelQuery(ctx)}`
   ),
};