@saber2pr/ai-agent 0.0.3 → 0.0.5

package/lib/agent.d.ts

@@ -4,28 +4,26 @@ export default class McpAgent {
     private clients;
     private allTools;
     private messages;
+    private engine;
+    private encoder;
     constructor();
     /**
-     * 1. API Configuration Management
-     * Checks for existing config or prompts user for input.
+     * 计算当前消息列表的总 Token 消耗
+     * 兼容多模态内容 (Content Parts) 和 工具调用 (Tool Calls)
      */
-    private ensureApiConfig;
+    private calculateTokens;
     /**
-     * 2. Load MCP server configs from common IDE paths
+     * 核心功能：内置代码分析工具
+     * 这里的逻辑直接调用 PromptEngine，不走网络请求，效率极高
      */
+    private registerBuiltinTools;
+    private ensureApiConfig;
     private loadMcpConfigs;
-    /**
-     * 3. Initialization
-     * Validates API credentials and connects to MCP servers.
-     */
     init(): Promise<void>;
-    /**
-     * 4. Core Chat Logic
-     * Handles user input and recursive tool calls.
-     */
     private processChat;
     /**
-     * 5. Start the Interactive Shell
+     * 简易 Loading 动画辅助函数
      */
+    private showLoading;
     start(): Promise<void>;
 }

package/lib/agent.js

@@ -36,124 +36,252 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
+const openai_1 = __importDefault(require("openai"));
 const index_js_1 = require("@modelcontextprotocol/sdk/client/index.js");
 const stdio_js_1 = require("@modelcontextprotocol/sdk/client/stdio.js");
-const openai_1 = __importDefault(require("openai"));
 const fs_1 = __importDefault(require("fs"));
 const path_1 = __importDefault(require("path"));
 const os_1 = __importDefault(require("os"));
 const readline = __importStar(require("readline"));
+const ts_context_mcp_1 = require("@saber2pr/ts-context-mcp"); // 引入我们的核心引擎
+const js_tiktoken_1 = require("js-tiktoken");
 const CONFIG_FILE = path_1.default.join(os_1.default.homedir(), ".saber2pr-agent.json");
-// --- Core Class ---
 class McpAgent {
     constructor() {
         this.modelName = "";
         this.clients = [];
         this.allTools = [];
         this.messages = [];
+        this.encoder = (0, js_tiktoken_1.getEncoding)("cl100k_base");
+        // 默认以当前工作目录为分析目标
+        this.engine = new ts_context_mcp_1.PromptEngine(process.cwd());
         this.messages.push({
             role: "system",
-            content: "You are a powerful local assistant. You can access local tools provided by the user via the MCP protocol. Please answer questions by combining tool outputs and context.",
+            content: `你是一个专业的 AI 代码架构师。
+你可以访问本地文件系统并利用 AST (抽象语法树) 技术分析代码。
+你的核心目标是提供准确的代码结构、依赖关系和逻辑分析。
+请先使用 get_repo_map 查看项目整体代码结构。
+请优先使用 read_skeleton 查看结构，只有在必要时才使用 read_full_code 或 get_method_body。`,
         });
+        // 初始化内置工具
+        this.registerBuiltinTools();
     }
     /**
-     * 1. API Configuration Management
-     * Checks for existing config or prompts user for input.
+     * 计算当前消息列表的总 Token 消耗
+     * 兼容多模态内容 (Content Parts) 和 工具调用 (Tool Calls)
      */
-    async ensureApiConfig() {
-        if (fs_1.default.existsSync(CONFIG_FILE)) {
-            try {
-                const config = JSON.parse(fs_1.default.readFileSync(CONFIG_FILE, "utf-8"));
-                if (config.baseURL && config.apiKey && config.model) {
-                    return config;
+    calculateTokens() {
+        let total = 0;
+        for (const msg of this.messages) {
+            // 1. 处理消息内容 (Content)
+            if (msg.content) {
+                if (typeof msg.content === "string") {
+                    // 普通文本消息
+                    total += this.encoder.encode(msg.content).length;
+                }
+                else if (Array.isArray(msg.content)) {
+                    // 多模态/复合内容消息 (ChatCompletionContentPart[])
+                    for (const part of msg.content) {
+                        if (part.type === "text" && "text" in part) {
+                            total += this.encoder.encode(part.text || "").length;
+                        }
+                        // 注意：图片 (image_url) 的 Token 计算通常基于分辨率，tiktoken 无法计算
+                    }
                 }
             }
-            catch (e) {
-                console.error(`[Error] Failed to read ${CONFIG_FILE}, re-initializing...`);
+            // 2. 处理助手角色发出的工具调用请求 (Assistant Tool Calls)
+            // 这是为了统计 AI 发出的指令所占用的 Token
+            if (msg.role === "assistant" && msg.tool_calls) {
+                for (const call of msg.tool_calls) {
+                    if (call.type === "function") {
+                        // 统计函数名和参数字符串
+                        total += this.encoder.encode(call.function.name).length;
+                        total += this.encoder.encode(call.function.arguments).length;
+                    }
+                }
+            }
+            // 3. 处理工具返回的结果 (Tool Role)
+            // 在 processChat 中，我们确保了工具返回的 result 最终被转为了 string
+            if (msg.role === "tool" && typeof msg.content === "string") {
+                total += this.encoder.encode(msg.content).length;
             }
         }
-        const rl = readline.createInterface({ input: process.stdin, output: process.stdout });
-        const question = (query) => new Promise((resolve) => rl.question(query, resolve));
-        console.log("\n🔑 API Configuration not found. Please provide the following details:");
-        const baseURL = await question("? API Base URL: ");
-        const apiKey = await question("? API Key: ");
-        const model = await question("? Model Name: ");
-        if (!baseURL || !apiKey || !model) {
-            console.error("❌ Error: All fields (Base URL, API Key, Model Name) are required!");
-            process.exit(1);
+        return total;
+    }
+    /**
+     * 核心功能：内置代码分析工具
+     * 这里的逻辑直接调用 PromptEngine，不走网络请求，效率极高
+     */
+    registerBuiltinTools() {
+        const builtinTools = [
+            {
+                type: "function",
+                function: {
+                    name: "get_repo_map",
+                    description: "获取项目全局文件结构及导出清单，用于快速定位代码",
+                    parameters: { type: "object", properties: {} },
+                },
+                _handler: async () => {
+                    this.engine.refresh();
+                    return this.engine.getRepoMap();
+                },
+            },
+            {
+                type: "function",
+                function: {
+                    name: "analyze_deps",
+                    description: "分析指定文件的依赖关系，支持 tsconfig 路径别名解析",
+                    parameters: {
+                        type: "object",
+                        properties: {
+                            filePath: { type: "string", description: "文件相对路径" },
+                        },
+                        required: ["filePath"],
+                    },
+                },
+                _handler: async ({ filePath }) => this.engine.getDeps(filePath),
+            },
+            {
+                type: "function",
+                function: {
+                    name: "read_skeleton",
+                    description: "提取文件的结构定义（接口、类、方法签名），忽略具体实现以节省 Token",
+                    parameters: {
+                        type: "object",
+                        properties: {
+                            filePath: { type: "string", description: "文件相对路径" },
+                        },
+                        required: ["filePath"],
+                    },
+                },
+                _handler: async (args) => {
+                    // 1. 严格路径守卫：防止 undefined 或空字符串进入 path 模块
+                    const pathArg = args === null || args === void 0 ? void 0 : args.filePath;
+                    if (typeof pathArg !== 'string' || pathArg.trim() === '') {
+                        return `Error: 参数 'filePath' 无效。收到的是: ${JSON.stringify(pathArg)}`;
+                    }
+                    try {
+                        // 2. 刷新引擎状态，确保分析的是最新的文件内容
+                        this.engine.refresh();
+                        // 3. 执行获取
+                        const result = this.engine.getSkeleton(pathArg);
+                        // 4. 空值回退：防止 getSkeleton 返回 null 导致后续统计 Token 时崩溃
+                        return result || `// Warning: 文件 ${pathArg} 存在但未找到任何可提取的结构。`;
+                    }
+                    catch (error) {
+                        // 5. 捕获 AST 级别的 match 错误
+                        return `Error: 解析文件 ${pathArg} 时发生内部错误: ${error.message}`;
+                    }
+                },
+            },
+            {
+                type: "function",
+                function: {
+                    name: "get_method_body",
+                    description: "获取指定文件内某个方法或函数的完整实现代码",
+                    parameters: {
+                        type: "object",
+                        properties: {
+                            filePath: { type: "string", description: "文件路径" },
+                            methodName: { type: "string", description: "方法名或函数名" },
+                        },
+                        required: ["filePath", "methodName"],
+                    },
+                },
+                _handler: async ({ filePath, methodName }) => this.engine.getMethodImplementation(filePath, methodName),
+            },
+            {
+                type: "function",
+                function: {
+                    name: "read_full_code",
+                    description: "读取指定文件的完整源代码内容。当需要分析具体实现逻辑或查找硬编码字符串时使用。",
+                    parameters: {
+                        type: "object",
+                        properties: {
+                            filePath: { type: "string", description: "文件相对路径" },
+                        },
+                        required: ["filePath"],
+                    },
+                },
+                // 核心实现：直接利用 fs 读取
+                _handler: async ({ filePath }) => {
+                    try {
+                        if (typeof filePath !== 'string' || !filePath) {
+                            return "Error: filePath 不能为空";
+                        }
+                        // 拼合绝对路径
+                        const fullPath = path_1.default.resolve(this.engine.getRootDir(), filePath);
+                        // 安全检查：防止 AI 尝试读取项目外的敏感文件
+                        if (!fullPath.startsWith(this.engine.getRootDir())) {
+                            return "Error: 权限拒绝，禁止访问项目目录外的文件。";
+                        }
+                        if (!fs_1.default.existsSync(fullPath)) {
+                            return `Error: 文件不存在: ${filePath}`;
+                        }
+                        return fs_1.default.readFileSync(fullPath, "utf-8");
+                    }
+                    catch (err) {
+                        return `Error: 读取文件失败: ${err.message}`;
+                    }
+                },
+            },
+        ];
+        this.allTools.push(...builtinTools);
+    }
+    // --- 初始化与环境准备 (API Config & MCP Servers) ---
+    async ensureApiConfig() {
+        if (fs_1.default.existsSync(CONFIG_FILE)) {
+            return JSON.parse(fs_1.default.readFileSync(CONFIG_FILE, "utf-8"));
         }
-        const config = { baseURL, apiKey, model };
+        const rl = readline.createInterface({
+            input: process.stdin,
+            output: process.stdout,
+        });
+        const question = (q) => new Promise((res) => rl.question(q, res));
+        console.log("\n🔑 配置 API 凭据:");
+        const config = {
+            baseURL: await question("? API Base URL (如 https://api.openai.com/v1): "),
+            apiKey: await question("? API Key: "),
+            model: await question("? Model Name (如 gpt-4o): "),
+        };
         fs_1.default.writeFileSync(CONFIG_FILE, JSON.stringify(config, null, 2));
-        console.log(`✅ Configuration saved to ${CONFIG_FILE}\n`);
         rl.close();
         return config;
     }
-    /**
-     * 2. Load MCP server configs from common IDE paths
-     */
     loadMcpConfigs() {
-        const combinedConfig = { mcpServers: {} };
-        const configPaths = [
+        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"),
         ];
-        for (const p of configPaths) {
+        paths.forEach((p) => {
             if (fs_1.default.existsSync(p)) {
-                try {
-                    const content = JSON.parse(fs_1.default.readFileSync(p, "utf-8"));
-                    if (content.mcpServers) {
-                        combinedConfig.mcpServers = { ...combinedConfig.mcpServers, ...content.mcpServers };
-                        console.log(`[MCP] Config loaded from: ${p}`);
-                    }
-                }
-                catch (e) {
-                    console.error(`[Error] Failed to parse MCP config ${p}:`, e);
-                }
+                const content = JSON.parse(fs_1.default.readFileSync(p, "utf-8"));
+                Object.assign(combined.mcpServers, content.mcpServers);
             }
-        }
-        return combinedConfig;
+        });
+        return combined;
     }
-    /**
-     * 3. Initialization
-     * Validates API credentials and connects to MCP servers.
-     */
     async init() {
-        // A. Setup & Validate OpenAI
         const apiConfig = await this.ensureApiConfig();
         this.openai = new openai_1.default({
             baseURL: apiConfig.baseURL,
             apiKey: apiConfig.apiKey,
         });
         this.modelName = apiConfig.model;
-        console.log("🔍 Validating API configuration...");
-        try {
-            // Perform a lightweight check to verify URL and Key
-            await this.openai.models.list();
-            console.log("✅ API validation successful.");
-        }
-        catch (e) {
-            console.error("\n❌ API Connection Failed!");
-            console.error(`Reason: ${e.message}`);
-            console.log(`\nSuggestion: If you made a mistake, please delete or edit: ${CONFIG_FILE}`);
-            process.exit(1);
-        }
-        // B. Setup MCP Clients
+        // 链接外部 MCP Server (如 Google Search, Filesystem 等)
         const mcpConfig = this.loadMcpConfigs();
-        const serverEntries = Object.entries(mcpConfig.mcpServers);
-        if (serverEntries.length === 0) {
-            console.warn("⚠️ No MCP server configurations found.");
-        }
-        for (const [name, server] of serverEntries) {
+        for (const [name, server] of Object.entries(mcpConfig.mcpServers)) {
             try {
                 const transport = new stdio_js_1.StdioClientTransport({
                     command: server.command,
                     args: server.args || [],
-                    env: { ...process.env, ...(server.env || {}) },
+                    env: { ...process.env, ...server.env },
                 });
                 const client = new index_js_1.Client({ name, version: "1.0.0" }, { capabilities: {} });
                 await client.connect(transport);
                 const { tools } = await client.listTools();
-                const formatted = tools.map((t) => ({
+                this.allTools.push(...tools.map((t) => ({
                     type: "function",
                     function: {
                         name: `${name}__${t.name}`,
@@ -162,93 +290,108 @@ class McpAgent {
                     },
                     _originalName: t.name,
                     _client: client,
-                }));
-                this.allTools.push(...formatted);
-                this.clients.push(client);
-                console.log(`✅ [${name}] Connected, loaded ${tools.length} tools`);
+                })));
+                console.log(`✅ [${name}] 加载成功`);
             }
             catch (e) {
-                console.error(`❌ [${name}] Failed to start:`, e);
+                console.error(`❌ [${name}] 启动失败`);
             }
         }
     }
-    /**
-     * 4. Core Chat Logic
-     * Handles user input and recursive tool calls.
-     */
     async processChat(userInput) {
-        this.messages.push({ role: "user", content: userInput });
-        let isThinking = true;
-        while (isThinking) {
-            const apiTools = this.allTools.map(({ _originalName, _client, ...rest }) => rest);
-            const response = await this.openai.chat.completions.create({
-                model: this.modelName,
-                messages: this.messages,
-                tools: apiTools.length > 0 ? apiTools : undefined,
-                tool_choice: "auto",
-            });
+        var _a;
+        this.messages.push({ role: 'user', content: userInput });
+        while (true) {
+            // 打印当前上下文的累计 Token
+            const currentInputTokens = this.calculateTokens();
+            console.log(`\n📊 当前上下文累计: ${currentInputTokens} tokens`);
+            const stopLoading = this.showLoading("🤖 Agent 正在思考...");
+            let response;
+            try {
+                response = await this.openai.chat.completions.create({
+                    model: this.modelName,
+                    messages: this.messages,
+                    tools: this.allTools.map(({ _handler, _client, _originalName, ...rest }) => rest),
+                    tool_choice: 'auto'
+                });
+            }
+            finally {
+                stopLoading();
+            }
             const message = response.choices[0].message;
-            // If no more tool calls, exit loop and show final response
-            if (!message.tool_calls || message.tool_calls.length === 0) {
-                this.messages.push(message);
+            this.messages.push(message);
+            // 计算本次 AI 回复生成的 Token
+            const completionTokens = ((_a = response.usage) === null || _a === void 0 ? void 0 : _a.completion_tokens) ||
+                (message.content ? this.encoder.encode(message.content).length : 0);
+            console.log(`✨ AI 回复消耗: ${completionTokens} tokens`);
+            if (!message.tool_calls) {
                 console.log(`\n🤖 Agent: ${message.content}`);
-                isThinking = false;
                 break;
             }
-            // Handle tool calls requested by the model
-            this.messages.push(message);
-            console.log(`\n⚙️  Model requested ${message.tool_calls.length} tool calls...`);
-            for (const toolCall of message.tool_calls) {
-                const toolInfo = this.allTools.find((t) => t.function.name === toolCall.function.name);
-                if (toolInfo) {
-                    const args = JSON.parse(toolCall.function.arguments);
-                    console.log(`   - Executing: ${toolInfo.function.name}`);
-                    try {
-                        const result = await toolInfo._client.callTool({
-                            name: toolInfo._originalName,
-                            arguments: args,
-                        });
-                        this.messages.push({
-                            role: "tool",
-                            tool_call_id: toolCall.id,
-                            content: JSON.stringify(result.content),
-                        });
-                    }
-                    catch (error) {
-                        console.error(`   - Execution failed: ${error.message}`);
-                        this.messages.push({
-                            role: "tool",
-                            tool_call_id: toolCall.id,
-                            content: `Error: ${error.message}`,
-                        });
-                    }
+            console.log(`\n⚙️ 正在执行 ${message.tool_calls.length} 个操作...`);
+            for (const call of message.tool_calls) {
+                const tool = this.allTools.find(t => t.function.name === call.function.name);
+                const args = JSON.parse(call.function.arguments);
+                // 打印文件路径提示
+                if (args.filePath) {
+                    console.log(`   📂 正在查看文件: ${args.filePath}`);
+                }
+                console.log(`   🛠️  执行: ${call.function.name}`);
+                let result;
+                if (tool === null || tool === void 0 ? void 0 : tool._handler) {
+                    result = await tool._handler(args);
+                }
+                else if ((tool === null || tool === void 0 ? void 0 : tool._client) && tool._originalName) {
+                    const mcpRes = await tool._client.callTool({
+                        name: tool._originalName,
+                        arguments: args
+                    });
+                    result = mcpRes.content;
                 }
+                const resultContent = typeof result === "string" ? result : JSON.stringify(result);
+                // 打印工具返回结果的 Token 消耗
+                const toolResultTokens = this.encoder.encode(resultContent).length;
+                console.log(`   📝 工具输出: ${toolResultTokens} tokens`);
+                this.messages.push({
+                    role: 'tool',
+                    tool_call_id: call.id,
+                    content: resultContent
+                });
+                console.log(`   ✅ 完成: ${call.function.name}`);
             }
         }
     }
     /**
-     * 5. Start the Interactive Shell
+     * 简易 Loading 动画辅助函数
      */
+    showLoading(text) {
+        const chars = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
+        let i = 0;
+        const timer = setInterval(() => {
+            process.stdout.write(`\r${chars[i]} ${text}`);
+            i = (i + 1) % chars.length;
+        }, 80);
+        return () => {
+            clearInterval(timer);
+            process.stdout.write('\r\x1b[K'); // 清除当前行
+        };
+    }
     async start() {
         await this.init();
         const rl = readline.createInterface({
             input: process.stdin,
             output: process.stdout,
         });
-        console.log(`\n🚀 Agent Started (Model: ${this.modelName})! Type 'exit' to quit.`);
+        console.log(`\n🚀 代码助手已启动 (目标目录: ${this.engine.getRootDir()})`);
         const chatLoop = () => {
-            rl.question("\n👤 You: ", async (input) => {
-                if (input.toLowerCase() === "exit") {
-                    console.log("Goodbye!");
-                    rl.close();
+            rl.question("\n👤 你: ", async (input) => {
+                if (input.toLowerCase() === "exit")
                     process.exit(0);
-                }
                 try {
                     await this.processChat(input);
                 }
                 catch (err) {
-                    console.error("\n❌ System Error during chat:", err.message);
-                    console.log("Try checking your API configuration or network connection.");
+                    console.error("\n❌ 系统错误:", err.message);
                 }
                 chatLoop();
             });

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@saber2pr/ai-agent",
-  "version": "0.0.3",
+  "version": "0.0.5",
   "description": "AI Assistant CLI.",
   "author": "saber2pr",
   "license": "ISC",
@@ -22,6 +22,8 @@
   },
   "dependencies": {
     "@modelcontextprotocol/sdk": "^1.25.3",
+    "@saber2pr/ts-context-mcp": "^0.0.6",
+    "js-tiktoken": "^1.0.21",
     "openai": "^6.16.0"
   },
   "devDependencies": {