@saber2pr/ai-agent 0.0.64 → 0.0.65

package/lib/core/agent-graph.d.ts

@@ -1,6 +1,6 @@
-import { AIMessage, BaseMessage } from '@langchain/core/messages';
-import { GraphAgentOptions } from '../types/type';
-import { AgentGraphModel } from '../model/AgentGraphModel';
+import { AIMessage, BaseMessage } from "@langchain/core/messages";
+import { GraphAgentOptions } from "../types/type";
+import { AgentGraphModel } from "../model/AgentGraphModel";
 interface TokenUsage {
     total: number;
 }
@@ -30,7 +30,7 @@ export default class McpGraphAgent<T extends AgentGraphModel = any> {
      * 设置外部流式输出回调（如 VS Code Webview）。
      * 设置后，callModel 的流式输出将通过回调发送而非写入 stdout。
      */
-    setStreamOutput(callback: (chunk: string, type: 'think' | 'text') => void): void;
+    setStreamOutput(callback: (chunk: string, type: "think" | "text") => void): void;
     /**
      * 清除外部流式输出回调，恢复默认的 stdout 输出。
      */

package/lib/core/agent-graph.js

@@ -76,11 +76,11 @@ class McpGraphAgent {
         const cleanup = async () => {
             this.stopLoading();
             await this.closeMcpClients(); // 清理 MCP 连接
-            process.stdout.write('\u001B[?25h');
+            process.stdout.write("\u001B[?25h");
             process.exit(0);
         };
-        process.on('SIGINT', cleanup);
-        process.on('SIGTERM', cleanup);
+        process.on("SIGINT", cleanup);
+        process.on("SIGTERM", cleanup);
     }
     /**
      * 设置外部流式输出回调（如 VS Code Webview）。
@@ -96,7 +96,7 @@ class McpGraphAgent {
         this.streamOutputCallback = null;
     }
     printLoadedTools() {
-        console.log('\n🛠️  [Graph] 正在加载工具节点...');
+        console.log("\n🛠️  [Graph] 正在加载工具节点...");
         this.langchainTools.forEach((tool) => {
             // 工具名称
             console.log(`\n🧰 工具名: ${tool.name}`);
@@ -120,12 +120,12 @@ class McpGraphAgent {
     loadMcpConfigs() {
         const combined = { mcpServers: {} };
         const paths = [
-            path_1.default.join(os_1.default.homedir(), '.cursor', 'mcp.json'),
-            path_1.default.join(os_1.default.homedir(), '.vscode', 'mcp.json'),
+            path_1.default.join(os_1.default.homedir(), ".cursor", "mcp.json"),
+            path_1.default.join(os_1.default.homedir(), ".vscode", "mcp.json"),
         ];
-        paths.forEach(p => {
+        paths.forEach((p) => {
             if (fs_1.default.existsSync(p)) {
-                const content = JSON.parse(fs_1.default.readFileSync(p, 'utf-8'));
+                const content = JSON.parse(fs_1.default.readFileSync(p, "utf-8"));
                 Object.assign(combined.mcpServers, content.mcpServers);
             }
         });
@@ -142,13 +142,13 @@ class McpGraphAgent {
                     args: config.args,
                     env: { ...process.env, ...(config.env || {}) },
                 });
-                const client = new index_js_1.Client({ name: 'mcp-graph-client', version: '1.0.0' }, { capabilities: {} });
+                const client = new index_js_1.Client({ name: "mcp-graph-client", version: "1.0.0" }, { capabilities: {} });
                 await client.connect(transport);
                 this.mcpClients.push(client);
                 const { tools } = await client.listTools();
-                tools.forEach(tool => {
+                tools.forEach((tool) => {
                     mcpToolInfos.push({
-                        type: 'function',
+                        type: "function",
                         function: {
                             name: tool.name,
                             description: tool.description,
@@ -172,11 +172,20 @@ class McpGraphAgent {
         return mcpToolInfos;
     }
     async prepareTools() {
-        const builtinToolInfos = (0, builtin_1.createDefaultBuiltinTools)({ options: this.options });
+        const builtinToolInfos = (0, builtin_1.createDefaultBuiltinTools)({
+            options: this.options,
+        });
         const mcpToolInfos = await this.initMcpTools();
         // 合并内置、手动传入和 MCP 工具
-        const allToolInfos = [...builtinToolInfos, ...(this.options.tools || []), ...mcpToolInfos];
-        this.langchainTools = allToolInfos.map(t => (0, convertToLangChainTool_1.convertToLangChainTool)(t, { allTools: allToolInfos }));
+        let allToolInfos = [
+            ...builtinToolInfos,
+            ...(this.options.tools || []),
+            ...mcpToolInfos,
+        ];
+        if (this.options.filterTools) {
+            allToolInfos = allToolInfos.filter(this.options.filterTools);
+        }
+        this.langchainTools = allToolInfos.map((t) => (0, convertToLangChainTool_1.convertToLangChainTool)(t, { allTools: allToolInfos }));
         this.toolNode = new prebuilt_1.ToolNode(this.langchainTools);
         return {
             builtinToolInfos,
@@ -190,7 +199,6 @@ class McpGraphAgent {
         if (this.model && this.langchainTools.length > 0) {
             return this.getModel();
         }
-        ;
         // 1. 加载所有工具（含 MCP）
         const toolsInfo = await this.prepareTools();
         // 2. 初始化模型
@@ -213,17 +221,17 @@ class McpGraphAgent {
         this.mcpClients = [];
     }
     showLoading(text) {
-        const chars = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
+        const chars = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
         let i = 0;
-        process.stdout.write('\u001B[?25l');
+        process.stdout.write("\u001B[?25l");
         const timer = setInterval(() => {
             process.stdout.write(`\r\x1b[36m${chars[i]}\x1b[0m ${text}`);
             i = (i + 1) % chars.length;
         }, 80);
         return () => {
             clearInterval(timer);
-            process.stdout.write('\r\x1b[K');
-            process.stdout.write('\u001B[?25h');
+            process.stdout.write("\r\x1b[K");
+            process.stdout.write("\u001B[?25h");
         };
     }
     startLoading(text) {
@@ -260,29 +268,33 @@ class McpGraphAgent {
         let config = {};
         if (fs_1.default.existsSync(config_1.CONFIG_FILE)) {
             try {
-                config = JSON.parse(fs_1.default.readFileSync(config_1.CONFIG_FILE, 'utf-8'));
+                config = JSON.parse(fs_1.default.readFileSync(config_1.CONFIG_FILE, "utf-8"));
             }
             catch (e) { }
         }
         if (!config.baseURL || !config.apiKey) {
-            const rl = readline_1.default.createInterface({ input: process.stdin, output: process.stdout });
-            const question = (q) => new Promise(res => rl.question(q, res));
+            const rl = readline_1.default.createInterface({
+                input: process.stdin,
+                output: process.stdout,
+            });
+            const question = (q) => new Promise((res) => rl.question(q, res));
             config.baseURL = config.baseURL || (await question(`? API Base URL: `));
             config.apiKey = config.apiKey || (await question(`? API Key: `));
-            config.model = config.model || (await question(`? Model Name: `)) || 'gpt-4o';
+            config.model =
+                config.model || (await question(`? Model Name: `)) || "gpt-4o";
             fs_1.default.writeFileSync(config_1.CONFIG_FILE, JSON.stringify(config, null, 2));
             rl.close();
         }
         return config;
     }
-    async chat(query = '开始代码审计') {
+    async chat(query = "开始代码审计") {
         try {
             await this.ensureInitialized();
             const app = await this.createGraph();
             const graphStream = await app.stream({
                 messages: [new messages_1.HumanMessage(query)],
             }, {
-                configurable: { thread_id: 'auto_worker' },
+                configurable: { thread_id: "auto_worker" },
                 recursionLimit: this.recursionLimit,
                 debug: this.verbose,
             });
@@ -290,7 +302,7 @@ class McpGraphAgent {
                 this.renderOutput(output, this.streamEnabled);
         }
         catch (error) {
-            console.error('\n❌ Chat 过程中发生错误:', error);
+            console.error("\n❌ Chat 过程中发生错误:", error);
         }
         finally {
             await this.closeMcpClients();
@@ -300,7 +312,7 @@ class McpGraphAgent {
      * 流式执行单次查询（编程式 API）。
      * 无论 options.stream 是否开启，此方法始终以流式方式输出。
      */
-    async stream(query = '开始代码审计') {
+    async stream(query = "开始代码审计") {
         const prevStream = this.streamEnabled;
         this.streamEnabled = true;
         try {
@@ -309,7 +321,7 @@ class McpGraphAgent {
             const graphStream = await app.stream({
                 messages: [new messages_1.HumanMessage(query)],
             }, {
-                configurable: { thread_id: 'stream_worker' },
+                configurable: { thread_id: "stream_worker" },
                 recursionLimit: this.recursionLimit,
                 debug: this.verbose,
             });
@@ -317,7 +329,7 @@ class McpGraphAgent {
                 this.renderOutput(output, true);
         }
         catch (error) {
-            console.error('\n❌ Stream 过程中发生错误:', error);
+            console.error("\n❌ Stream 过程中发生错误:", error);
         }
         finally {
             this.streamEnabled = prevStream;
@@ -327,21 +339,24 @@ class McpGraphAgent {
     async start() {
         await this.ensureInitialized();
         const app = await this.createGraph();
-        const rl = readline_1.default.createInterface({ input: process.stdin, output: process.stdout });
-        rl.on('SIGINT', () => {
+        const rl = readline_1.default.createInterface({
+            input: process.stdin,
+            output: process.stdout,
+        });
+        rl.on("SIGINT", () => {
             this.stopLoading();
             rl.close();
-            process.stdout.write('\u001B[?25h');
+            process.stdout.write("\u001B[?25h");
             process.exit(0);
         });
         const ask = () => {
-            rl.question('> ', async (input) => {
-                if (input.toLowerCase() === 'exit') {
+            rl.question("> ", async (input) => {
+                if (input.toLowerCase() === "exit") {
                     rl.close();
                     return;
                 }
                 const graphStream = await app.stream({ messages: [new messages_1.HumanMessage(input)] }, {
-                    configurable: { thread_id: 'session' },
+                    configurable: { thread_id: "session" },
                     recursionLimit: this.recursionLimit,
                     debug: this.verbose,
                 });
@@ -358,7 +373,9 @@ class McpGraphAgent {
         // ✅ 打印工具执行结果（tools 节点的输出）
         const toolsNode = output.tools;
         if (toolsNode && toolsNode.messages) {
-            const toolMessages = Array.isArray(toolsNode.messages) ? toolsNode.messages : [];
+            const toolMessages = Array.isArray(toolsNode.messages)
+                ? toolsNode.messages
+                : [];
             // 获取最近的 AI 消息以匹配 tool_call_id
             const lastAiMsg = agentNode?.messages?.[agentNode.messages.length - 1];
             const toolCallMap = new Map();
@@ -372,10 +389,12 @@ class McpGraphAgent {
                 // ToolMessage 有 tool_call_id 字段
                 const toolCallId = msg.tool_call_id || msg.id;
                 if (toolCallId) {
-                    const toolName = toolCallMap.get(toolCallId) || msg.name || 'unknown';
-                    const content = typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content);
+                    const toolName = toolCallMap.get(toolCallId) || msg.name || "unknown";
+                    const content = typeof msg.content === "string"
+                        ? msg.content
+                        : JSON.stringify(msg.content);
                     // 如果内容太长，截断显示
-                    const displayContent = content.length > 500 ? content.substring(0, 500) + '...' : content;
+                    const displayContent = content.length > 500 ? content.substring(0, 500) + "..." : content;
                     console.log(`✅ [工具结果] ${toolName}: ${displayContent}`);
                 }
             });
@@ -402,7 +421,7 @@ class McpGraphAgent {
             }
             // 4. 打印工具调用情况
             if (msg.tool_calls?.length) {
-                msg.tool_calls.forEach(call => {
+                msg.tool_calls.forEach((call) => {
                     console.log(`🛠️ [调用工具]: ${call.name} 📦 参数: ${JSON.stringify(call.args)}`);
                 });
             }
@@ -412,9 +431,9 @@ class McpGraphAgent {
         const recentToolCalls = this.getRecentToolCalls(state.messages);
         const recentToolCallsStr = recentToolCalls.length > 0
             ? `\n\n⚠️ 最近调用的工具（避免重复调用相同工具和参数）：\n${recentToolCalls
-                .map(tc => `  - ${tc.name}(${JSON.stringify(tc.args)})`)
-                .join('\n')}`
-            : '';
+                .map((tc) => `  - ${tc.name}(${JSON.stringify(tc.args)})`)
+                .join("\n")}`
+            : "";
         // 1. 构建当前的系统提示词模板
         const systemPromptTemplate = `
 ${(0, getSystemPromptTemplate_1.getSystemPromptTemplate)(this.targetDir)}
@@ -425,32 +444,32 @@ ${(0, getSystemPromptTemplate_1.getSystemPromptTemplate)(this.targetDir)}
         // ✅ 检查 options 中的 alwaysSystem 参数 (默认为 true 或根据你的需求设置)
         // 如果不希望每次都携带（即只在首轮携带），则过滤掉历史消息里的 SystemMessage
         if (this.options.alwaysSystem === false) {
-            inputMessages = state.messages.filter(msg => msg._getType() !== 'system');
+            inputMessages = state.messages.filter((msg) => msg._getType() !== "system");
         }
         else {
             // 默认模式：保持干净，由 PromptTemplate 重新生成最新的 System 状态
-            inputMessages = state.messages.filter(msg => msg._getType() !== 'system');
+            inputMessages = state.messages.filter((msg) => msg._getType() !== "system");
         }
         const prompt = prompts_1.ChatPromptTemplate.fromMessages([
-            ['system', systemPromptTemplate],
-            new prompts_1.MessagesPlaceholder('messages'),
+            ["system", systemPromptTemplate],
+            new prompts_1.MessagesPlaceholder("messages"),
         ]);
-        this.startLoading('AI 正在分析并思考中');
+        this.startLoading("AI 正在分析并思考中");
         try {
             const promptParams = {
                 messages: inputMessages,
                 recentToolCalls: recentToolCallsStr,
-                extraPrompt: this.options.extraSystemPrompt || '',
+                extraPrompt: this.options.extraSystemPrompt || "",
             };
             if (this.streamEnabled) {
                 // ✅ 流式模式：通过 AgentGraphModel.streamGenerate 进行流式输出
                 const formattedMessages = await prompt.formatMessages(promptParams);
                 this.stopLoading();
                 // --- 流式 <think> 标签实时过滤 + 流式打印思考内容 ---
-                const THINK_OPEN = '<think>';
-                const THINK_CLOSE = '</think>';
+                const THINK_OPEN = "<think>";
+                const THINK_CLOSE = "</think>";
                 let inThink = false;
-                let textBuffer = '';
+                let textBuffer = "";
                 let aiHeaderPrinted = false;
                 let thinkHeaderPrinted = false;
                 const hasExternalHandler = !!this.streamOutputCallback;
@@ -458,11 +477,11 @@ ${(0, getSystemPromptTemplate_1.getSystemPromptTemplate)(this.targetDir)}
                     if (!text)
                         return;
                     if (hasExternalHandler) {
-                        this.streamOutputCallback(text, 'text');
+                        this.streamOutputCallback(text, "text");
                     }
                     else {
                         if (!aiHeaderPrinted) {
-                            process.stdout.write('🤖 [AI]: ');
+                            process.stdout.write("🤖 [AI]: ");
                             aiHeaderPrinted = true;
                         }
                         process.stdout.write(text);
@@ -472,11 +491,11 @@ ${(0, getSystemPromptTemplate_1.getSystemPromptTemplate)(this.targetDir)}
                     if (!text)
                         return;
                     if (hasExternalHandler) {
-                        this.streamOutputCallback(text, 'think');
+                        this.streamOutputCallback(text, "think");
                     }
                     else {
                         if (!thinkHeaderPrinted) {
-                            process.stdout.write('\x1b[2m🧠 [思考]: ');
+                            process.stdout.write("\x1b[2m🧠 [思考]: ");
                             thinkHeaderPrinted = true;
                         }
                         process.stdout.write(text);
@@ -496,7 +515,7 @@ ${(0, getSystemPromptTemplate_1.getSystemPromptTemplate)(this.targetDir)}
                                 inThink = false;
                                 // 思考块结束：stdout 模式下换行 + 重置样式
                                 if (!hasExternalHandler && thinkHeaderPrinted) {
-                                    process.stdout.write('\x1b[0m\n');
+                                    process.stdout.write("\x1b[0m\n");
                                     thinkHeaderPrinted = false;
                                 }
                             }
@@ -539,10 +558,10 @@ ${(0, getSystemPromptTemplate_1.getSystemPromptTemplate)(this.targetDir)}
                 // 收尾：stdout 模式下关闭样式和换行
                 if (!hasExternalHandler) {
                     if (thinkHeaderPrinted) {
-                        process.stdout.write('\x1b[0m\n');
+                        process.stdout.write("\x1b[0m\n");
                     }
                     if (aiHeaderPrinted)
-                        process.stdout.write('\n');
+                        process.stdout.write("\n");
                 }
                 const aiMsg = result.generations[0].message;
                 const meta = aiMsg.response_metadata || {};
@@ -591,57 +610,49 @@ ${(0, getSystemPromptTemplate_1.getSystemPromptTemplate)(this.targetDir)}
         const totalTokens = state.tokenUsage?.total || 0;
         const totalMs = state.totalDuration || 0;
         if (totalTokens > 0 || totalMs > 0) {
-            console.log('\n' + '═'.repeat(50));
+            console.log("\n" + "═".repeat(50));
             console.log(`🏁 \x1b[32;1m[审计任务全量结算]\x1b[0m`);
             console.log(`   - 累计消耗总额: \x1b[33m${totalTokens}\x1b[0m Tokens`);
             console.log(`   - 累计执行耗时: \x1b[36m${(totalMs / 1000).toFixed(2)}\x1b[0m s`);
-            console.log('═'.repeat(50) + '\n');
+            console.log("═".repeat(50) + "\n");
         }
     }
     async createGraph() {
         const workflow = new langgraph_1.StateGraph(AgentState)
-            .addNode('agent', state => this.callModel(state))
-            .addNode('tools', this.toolNode)
-            .addNode('interrupt', state => {
+            .addNode("agent", (state) => this.callModel(state))
+            .addNode("tools", this.toolNode)
+            .addNode("interrupt", (state) => {
             return {
-                messages: [new messages_1.AIMessage({
-                        content: `Final Answer:\n💡 **嘿，我们已经聊了很多了！**
-
-为了保证服务质量，我通常会在 50 步操作后停下来和您确认一下。这样可以避免我产生“幻觉”或者做一些无用功。
-
-**任务还没完成吗？** 没关系，只要您回复“继续”，我马上就回来接着干！或者您有什么新的想法要告诉我？`
-                    })]
+                messages: [
+                    new messages_1.AIMessage({
+                        content: `Final Answer:\n💡 ⚠️ 抱歉，当前任务消耗的 Token 已达上限 ${this.options.maxTokens}。为了防止无限循环，我已停止执行。你可以尝试分步骤指令，继续发送指令，直到任务完成。`,
+                    }),
+                ],
             };
         })
-            .addEdge(langgraph_1.START, 'agent')
-            .addConditionalEdges('agent', state => {
+            .addEdge(langgraph_1.START, "agent")
+            .addConditionalEdges("agent", (state) => {
             const { messages } = state;
             const lastMsg = messages[messages.length - 1];
-            const content = lastMsg.content || '';
+            const content = lastMsg.content || "";
             // 🛑 新增：全局 Token 熔断保护
             // 如果已消耗 Token 超过了 options 中设置的 maxTokens (假设是总限额)
-            if (this.options.maxTokens && state.tokenUsage.total >= this.options.maxTokens) {
-                console.warn('⚠️ [警告] 已达到最大 Token 限制，强制结束任务。');
-                state.messages.push(new messages_1.AIMessage({
-                    content: `⚠️ 抱歉，当前任务消耗的 Token 已达上限 ${this.options.maxTokens}。为了防止无限循环，我已停止执行。你可以尝试分步骤指令，继续发送指令，直到任务完成。`
-                }));
-                this.printFinalSummary(state);
-                return langgraph_1.END;
+            if (this.options.maxTokens &&
+                state.tokenUsage.total >= this.options.maxTokens) {
+                return "interrupt";
             }
             // 1. 如果 AI 想要调用工具，去 tools 节点
             if (lastMsg.tool_calls && lastMsg.tool_calls.length > 0) {
-                return 'tools';
+                return "tools";
             }
-            const isFinalAnswer = content.includes('Final Answer');
+            const isFinalAnswer = content.includes("Final Answer");
             if (isFinalAnswer) {
                 this.printFinalSummary(state);
                 return langgraph_1.END;
             }
-            if (messages.length > 50) {
-                return 'interrupt';
-            }
             return langgraph_1.END;
-        }).addEdge('tools', 'agent');
+        })
+            .addEdge("tools", "agent");
         return workflow.compile({ checkpointer: this.checkpointer });
     }
 }

package/lib/types/type.d.ts

@@ -1,13 +1,13 @@
-import { z } from 'zod';
-import { Client } from '@modelcontextprotocol/sdk/client/index';
-import { AgentGraphModel } from '../model/AgentGraphModel';
+import { z } from "zod";
+import { Client } from "@modelcontextprotocol/sdk/client/index";
+import { AgentGraphModel } from "../model/AgentGraphModel";
 export interface ApiConfig {
     baseURL: string;
     apiKey: string;
     model: string;
 }
 export interface ToolInfo {
-    type: 'function';
+    type: "function";
     function: {
         name: string;
         description?: string;
@@ -42,6 +42,7 @@ export interface GraphAgentOptions<T extends AgentGraphModel = any> extends Agen
     recursionLimit?: number;
     /** 是否启用流式输出，默认 false */
     stream?: boolean;
+    filterTools?: (tool: ToolInfo) => boolean;
 }
 export interface CreateAgentOptions {
     apiKey: string;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@saber2pr/ai-agent",
-  "version": "0.0.64",
+  "version": "0.0.65",
   "description": "AI Assistant CLI.",
   "author": "saber2pr",
   "license": "ISC",