fix: 增加Deepwiki 的一次性链接模式 (#28)

* fix: 增加Deepwiki 的一次性链接模式

* fix: 修复 CR 意见

* fix: 修复 CR 意见

Author: Alexzjt

__tests__/utils/deepwiki.test.ts

@@ -1,26 +1,91 @@
-import { describe, it, expect } from 'vitest';
-import { queryDeepWiki } from '../../src/utils/deepwiki';
+import { describe, it, expect, afterAll } from 'vitest';
+import {
+  queryDeepWiki,
+  adaptedQueryDeepWiki,
+  closeDeepWikiConnection,
+  getRepoName,
+} from '../../src/utils/deepwiki';
+
+afterAll(async () => {
+  await closeDeepWikiConnection();
+});
+
+describe('getRepoName', () => {
+  it('should map adc to ant-design/ant-design-charts', () => {
+    expect(getRepoName('adc')).toBe('ant-design/ant-design-charts');
+  });
+
+  it('should map normal library to antvis/<UPPER>', () => {
+    expect(getRepoName('g2')).toBe('antvis/G2');
+    expect(getRepoName('s2')).toBe('antvis/S2');
+    expect(getRepoName('l7')).toBe('antvis/L7');
+  });
+});
 
 describe('DeepWiki Integration Test', () => {
   it('should get real response from DeepWiki for G2 question', async () => {
-    // Set timeout to 120s as this is a real network request
-
-    try {
-      const result = await queryDeepWiki({
-        repoName: 'antvis/G2',
-        question: '如何调整折线图两端的间隔',
-      });
-
-      console.log('DeepWiki Answer:', result);
-
-      expect(result).toBeDefined();
-      expect(typeof result).toBe('string');
-      expect(result.length).toBeGreaterThan(10);
-      // Basic check to see if it returned something relevant or at least not an error string
-      expect(result).not.toMatch(/^Error/);
-    } catch (error) {
-      console.error('DeepWiki Query Failed:', error);
-      throw error;
-    }
+    const result = await queryDeepWiki({
+      repoName: 'antvis/G2',
+      question: '如何调整折线图两端的间隔',
+    });
+
+    console.log('DeepWiki Answer (keep-alive):', result.slice(0, 200));
+
+    expect(result).toBeDefined();
+    expect(typeof result).toBe('string');
+    expect(result.length).toBeGreaterThan(10);
+    expect(result).not.toMatch(/^Error/);
+    // 确保尾部 wiki 推荐链接已被清理
+    expect(result).not.toMatch(/Wiki pages you might want to explore/i);
+    expect(result).not.toMatch(/View this search on DeepWiki/i);
+  }, 120000);
+
+  it('should work with close-after-query mode', async () => {
+    const result = await queryDeepWiki({
+      repoName: 'antvis/G2',
+      question: 'What is G2?',
+      connectionMode: 'close-after-query',
+    });
+
+    console.log('DeepWiki Answer (close-after-query):', result.slice(0, 200));
+
+    expect(result).toBeDefined();
+    expect(typeof result).toBe('string');
+    expect(result.length).toBeGreaterThan(10);
+  }, 120000);
+
+  it('should auto-resolve short repoName via getRepoName', async () => {
+    const result = await queryDeepWiki({
+      repoName: 'g2',
+      question: 'What charts does G2 support?',
+    });
+
+    expect(result).toBeDefined();
+    expect(result.length).toBeGreaterThan(10);
+  }, 120000);
+});
+
+describe('adaptedQueryDeepWiki', () => {
+  it('should return { documentation } on success', async () => {
+    const result = await adaptedQueryDeepWiki({
+      repoName: 'antvis/G2',
+      question: 'What is G2?',
+    });
+
+    expect(result.documentation).toBeDefined();
+    expect(typeof result.documentation).toBe('string');
+    expect(result).not.toHaveProperty('error');
+  }, 120000);
+
+  it('should return { documentation: null, error } on failure', async () => {
+    const result = await adaptedQueryDeepWiki({
+      repoName: 'nonexistent/repo-that-does-not-exist-12345',
+      question: 'anything',
+    });
+
+    // DeepWiki 对不存在的仓库可能返回空或报错
+    // adaptedQueryDeepWiki 应该把异常吞掉，返回 error 字段
+    expect(result.documentation).toBeNull();
+    expect(typeof result.error).toBe('string');
   }, 120000);
 });

src/utils/deepwiki.ts

@@ -4,6 +4,19 @@ import EventSource from 'eventsource';
 import { logger } from './logger';
 import { AntVLibrary } from '../types';
 
+export type DeepWikiConnectionMode = 'keep-alive' | 'close-after-query';
+
+type DeepWikiQueryParams = {
+  repoName: string;
+  question: string;
+  connectionMode?: DeepWikiConnectionMode;
+};
+
+type McpConnection = {
+  client: Client;
+  transport: StreamableHTTPClientTransport;
+};
+
 // ---------------------------------------------------------
 // 1. 全局环境配置
 // ---------------------------------------------------------
@@ -21,6 +34,77 @@ let _client: Client | null = null;
 let _transport: StreamableHTTPClientTransport | null = null;
 let _connectingPromise: Promise<Client> | null = null;
 
+function createMcpConnection(): McpConnection {
+  const transport = new StreamableHTTPClientTransport(
+    new URL('https://mcp.deepwiki.com/mcp'),
+  );
+
+  const client = new Client(
+    {
+      name: 'deepwiki-node-client',
+      version: '1.0.0',
+    },
+    { capabilities: {} },
+  );
+
+  return { client, transport };
+}
+
+async function connectMcpConnection(params: {
+  client: Client;
+  transport: StreamableHTTPClientTransport;
+  isShared?: boolean;
+}) {
+  const { client, transport, isShared = false } = params;
+
+  transport.onerror = (err) => {
+    logger.error('DeepWiki MCP Transport Error:', err);
+    if (isShared) {
+      resetClient();
+    }
+  };
+
+  transport.onclose = () => {
+    logger.info('DeepWiki MCP Connection Closed');
+    if (isShared) {
+      resetClient();
+    }
+  };
+
+  await client.connect(transport);
+}
+
+async function closeMcpConnection(params: {
+  client?: Client | null;
+  transport?: StreamableHTTPClientTransport | null;
+  isShared?: boolean;
+}) {
+  const { client, transport, isShared = false } = params;
+
+  if (!client && !transport) {
+    return;
+  }
+
+  try {
+    if (client) {
+      await client.close();
+    } else {
+      await transport?.close();
+    }
+  } catch (error) {
+    logger.error('DeepWiki MCP Client Close Error:', error);
+    try {
+      await transport?.close();
+    } catch (transportError) {
+      logger.error('DeepWiki MCP Transport Close Error:', transportError);
+    }
+  } finally {
+    if (isShared) {
+      resetClient();
+    }
+  }
+}
+
 /**
  * 获取或初始化 MCP 客户端 (单例模式)
  * 包含防止并发重复连接的锁机制
@@ -40,42 +124,21 @@ async function getMcpClient(): Promise<Client> {
 
   // 开始初始化连接
   _connectingPromise = (async () => {
+    let connection: McpConnection | null = null;
     try {
       logger.info('DeepWiki MCP: 初始化连接...');
-
-      const transport = new StreamableHTTPClientTransport(
-        new URL('https://mcp.deepwiki.com/mcp'),
-      );
-
-      const client = new Client(
-        {
-          name: 'deepwiki-node-client',
-          version: '1.0.0',
-        },
-        { capabilities: {} },
-      );
-
-      // 错误处理：监听传输层关闭或错误，以便下次重连
-      transport.onerror = (err) => {
-        logger.error('DeepWiki MCP Transport Error:', err);
-        resetClient(); // 出错时重置，下次调用会触发重连
-      };
-
-      transport.onclose = () => {
-        logger.info('DeepWiki MCP Connection Closed');
-        resetClient();
-      };
-
-      await client.connect(transport);
+      connection = createMcpConnection();
+      await connectMcpConnection({ ...connection, isShared: true });
 
       // 连接成功，赋值给模块级变量
-      _client = client;
-      _transport = transport;
+      _client = connection.client;
+      _transport = connection.transport;
       logger.info('DeepWiki MCP: 连接成功');
 
-      return client;
+      return connection.client;
     } catch (error) {
       logger.error('DeepWiki MCP: 连接失败', error);
+      await closeMcpConnection(connection ?? {});
       resetClient();
       throw error;
     } finally {
@@ -106,22 +169,36 @@ function resetClient() {
 export async function queryDeepWiki(_params: {
   repoName: string;
   question: string;
+  connectionMode?: DeepWikiConnectionMode;
 }): Promise<string> {
+  let localConnection: McpConnection | null = null;
   try {
-    const params = {
+    const params: DeepWikiQueryParams = {
       ..._params,
     };
     if (!params.repoName.includes('/')) {
       params.repoName = getRepoName(params.repoName as AntVLibrary);
     }
 
-    // 获取单例客户端 (自动处理连接)
-    const client = await getMcpClient();
+    let client: Client;
+    if (params.connectionMode === 'close-after-query') {
+      logger.info('DeepWiki MCP: 初始化一次性连接...');
+      localConnection = createMcpConnection();
+      await connectMcpConnection(localConnection);
+      client = localConnection.client;
+      logger.info('DeepWiki MCP: 一次性连接成功');
+    } else {
+      // 获取单例客户端 (自动处理连接)
+      client = await getMcpClient();
+    }
 
     // 调用工具
     const result: any = await client.callTool({
       name: 'ask_question',
-      arguments: params,
+      arguments: {
+        repoName: params.repoName,
+        question: params.question,
+      },
     });
 
     // ---------------------------------------------------------
@@ -134,6 +211,10 @@ export async function queryDeepWiki(_params: {
       .map((item: any) => item.text)
       .join('\n'); // 如果有多段文本，用换行符拼接
 
+    if (result.isError) {
+      throw new Error('DeepWiki Tool Error, Answer = ' + answer);
+    }
+
     const regex =
       /Wiki pages you might want to explore:|View this search on DeepWiki:/i;
     const splitIndex = answer.search(regex);
@@ -159,12 +240,18 @@ export async function queryDeepWiki(_params: {
       // resetClient();
     }
     throw error;
+  } finally {
+    if (localConnection) {
+      logger.info('DeepWiki MCP: 关闭一次性连接...');
+      await closeMcpConnection(localConnection);
+    }
   }
 }
 
 export async function adaptedQueryDeepWiki(_params: {
   repoName: string;
   question: string;
+  connectionMode?: DeepWikiConnectionMode;
 }) {
   try {
     const answer = await queryDeepWiki(_params);
@@ -182,16 +269,17 @@ export async function adaptedQueryDeepWiki(_params: {
  * 供 EggJS 应用在 app.beforeClose 时调用
  */
 export async function closeDeepWikiConnection() {
-  if (_transport) {
+  if (_connectingPromise) {
+    await _connectingPromise.catch(() => {});
+  }
+
+  if (_client || _transport) {
     logger.info('DeepWiki MCP: 关闭连接...');
-    // SDK 目前可能没有直接暴露 close 方法，通常关闭 transport 即可
-    // StreamableHTTPClientTransport 内部并没有显式的 close 方法暴露出来，
-    // 但我们可以将引用置空，让 GC 回收，或者依赖进程退出
-    // 如果 SSEClientTransport 实现了 close，则调用:
-    // await _transport.close();
-
-    // 目前 SDK 版本可以直接置空
-    resetClient();
+    await closeMcpConnection({
+      client: _client,
+      transport: _transport,
+      isShared: true,
+    });
   }
 }