@saber2pr/ai-agent 0.0.4 → 0.0.6

package/lib/agent.d.ts

@@ -1,3 +1,16 @@
+export interface CustomTool {
+    name: string;
+    description: string;
+    parameters: any;
+    handler: (args: any) => Promise<any>;
+}
+export interface AgentOptions {
+    targetDir?: string;
+    /** 自定义工具扩展 */
+    tools?: any[];
+    /** 注入到 System Prompt 中的额外指令/规则/上下文 */
+    extraSystemPrompt?: any;
+}
 export default class McpAgent {
     private openai;
     private modelName;
@@ -5,7 +18,15 @@ export default class McpAgent {
     private allTools;
     private messages;
     private engine;
-    constructor();
+    private encoder;
+    private extraTools;
+    constructor(options?: AgentOptions);
+    /**
+     * 计算当前消息列表的总 Token 消耗
+     * 兼容多模态内容 (Content Parts) 和 工具调用 (Tool Calls)
+     */
+    private calculateTokens;
+    private injectCustomTools;
     /**
      * 核心功能：内置代码分析工具
      * 这里的逻辑直接调用 PromptEngine，不走网络请求，效率极高
@@ -14,10 +35,16 @@ export default class McpAgent {
     private ensureApiConfig;
     private loadMcpConfigs;
     init(): Promise<void>;
+    private processChat;
     /**
-     * 核心交互循环 (Reasoning Loop)
-     * 允许 AI 连续调用工具来解决复杂代码问题
+     * 简易 Loading 动画辅助函数
      */
-    private processChat;
+    private showLoading;
+    /**
+     * 编程式对话入口
+     * @param input 用户指令
+     * @returns AI 的最终答复内容
+     */
+    chat(input: string): Promise<string>;
     start(): Promise<void>;
 }

package/lib/agent.js

@@ -44,24 +44,92 @@ 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");
 class McpAgent {
-    constructor() {
+    constructor(options) {
         this.modelName = "";
         this.clients = [];
         this.allTools = [];
         this.messages = [];
-        // 默认以当前工作目录为分析目标
-        this.engine = new ts_context_mcp_1.PromptEngine(process.cwd());
-        this.messages.push({
-            role: "system",
-            content: `你是一个专业的 AI 代码架构师。
+        this.encoder = (0, js_tiktoken_1.getEncoding)("cl100k_base");
+        this.extraTools = [];
+        this.engine = new ts_context_mcp_1.PromptEngine((options === null || options === void 0 ? void 0 : options.targetDir) || process.cwd());
+        this.extraTools = (options === null || options === void 0 ? void 0 : options.tools) || []; // 接收外部传入的工具
+        let baseSystemPrompt = `你是一个专业的 AI 代码架构师。
 你可以访问本地文件系统并利用 AST (抽象语法树) 技术分析代码。
 你的核心目标是提供准确的代码结构、依赖关系和逻辑分析。
-请优先使用 read_skeleton 查看结构，只有在必要时才使用 read_full_code 或 get_method_body。`,
+请先使用 get_repo_map 查看项目整体代码结构。
+请优先使用 read_skeleton 查看结构，只有在必要时才使用 read_full_code 或 get_method_body。`;
+        // 2. 拼接额外指令
+        if (options === null || options === void 0 ? void 0 : options.extraSystemPrompt) {
+            const extra = typeof options.extraSystemPrompt === 'string'
+                ? options.extraSystemPrompt
+                : JSON.stringify(options.extraSystemPrompt, null, 2);
+            baseSystemPrompt += `\n\n[额外执行指令]:\n${extra}`;
+        }
+        this.messages.push({
+            role: "system",
+            content: baseSystemPrompt,
         });
         // 初始化内置工具
         this.registerBuiltinTools();
+        this.injectCustomTools(); // 注入外部工具
+    }
+    /**
+     * 计算当前消息列表的总 Token 消耗
+     * 兼容多模态内容 (Content Parts) 和 工具调用 (Tool Calls)
+     */
+    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 无法计算
+                    }
+                }
+            }
+            // 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;
+            }
+        }
+        return total;
+    }
+    injectCustomTools() {
+        for (const tool of this.extraTools) {
+            this.allTools.push({
+                type: "function",
+                function: {
+                    name: tool.name,
+                    description: tool.description,
+                    parameters: tool.parameters,
+                },
+                _handler: tool.handler,
+            });
+        }
     }
     /**
      * 核心功能：内置代码分析工具
@@ -109,7 +177,25 @@ class McpAgent {
                         required: ["filePath"],
                     },
                 },
-                _handler: async ({ filePath }) => this.engine.getSkeleton(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",
@@ -127,6 +213,45 @@ class McpAgent {
                 },
                 _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}`;
+                        }
+                        const content = fs_1.default.readFileSync(fullPath, "utf-8");
+                        // 加上行号，AI 就能在 generate_review 里给出准确的 line 参数
+                        return content.split('\n')
+                            .map((line, i) => `${i + 1} | ${line}`)
+                            .join('\n');
+                    }
+                    catch (err) {
+                        return `Error: 读取文件失败: ${err.message}`;
+                    }
+                },
+            },
         ];
         this.allTools.push(...builtinTools);
     }
@@ -200,52 +325,105 @@ class McpAgent {
             }
         }
     }
-    /**
-     * 核心交互循环 (Reasoning Loop)
-     * 允许 AI 连续调用工具来解决复杂代码问题
-     */
     async processChat(userInput) {
-        this.messages.push({ role: "user", content: userInput });
+        var _a;
+        this.messages.push({ role: 'user', content: userInput });
         while (true) {
-            const 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",
-            });
+            // 打印当前上下文的累计 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;
             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}`);
                 break;
             }
-            console.log(`\n⚙️ 正在思考并执行 ${message.tool_calls.length} 个操作...`);
+            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 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) {
-                    // 执行内置 PromptEngine 工具
-                    result = await tool._handler(JSON.parse(call.function.arguments));
+                    result = await tool._handler(args);
                 }
                 else if ((tool === null || tool === void 0 ? void 0 : tool._client) && tool._originalName) {
-                    // 执行外部 MCP 工具
                     const mcpRes = await tool._client.callTool({
                         name: tool._originalName,
-                        arguments: JSON.parse(call.function.arguments),
+                        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",
+                    role: 'tool',
                     tool_call_id: call.id,
-                    content: typeof result === "string" ? result : JSON.stringify(result),
+                    content: resultContent
                 });
-                console.log(`   - 完成: ${call.function.name}`);
+                console.log(`   ✅ 完成: ${call.function.name}`);
             }
         }
     }
+    /**
+     * 简易 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'); // 清除当前行
+        };
+    }
+    /**
+     * 编程式对话入口
+     * @param input 用户指令
+     * @returns AI 的最终答复内容
+     */
+    async chat(input) {
+        if (!this.openai) {
+            await this.init();
+        }
+        // 调用现有的处理逻辑
+        // 假设你的 processChat 已经处理了所有的 tool_calls 循环
+        await this.processChat(input);
+        // 返回消息列表中的最后一条 AI 回复
+        const lastMsg = this.messages[this.messages.length - 1];
+        return lastMsg.role === 'assistant' ? lastMsg.content : '';
+    }
+    // 修改原来的 start 方法，使其内部也调用 chat
     async start() {
-        await this.init();
+        if (!this.openai) {
+            await this.init();
+        }
         const rl = readline.createInterface({
             input: process.stdin,
             output: process.stdout,
@@ -256,7 +434,8 @@ class McpAgent {
                 if (input.toLowerCase() === "exit")
                     process.exit(0);
                 try {
-                    await this.processChat(input);
+                    // 这里统一调用 chat 或核心逻辑
+                    await this.chat(input);
                 }
                 catch (err) {
                     console.error("\n❌ 系统错误:", err.message);

package/lib/cli.d.ts

@@ -0,0 +1,2 @@
+#!/usr/bin/env node
+export {};

package/lib/cli.js

@@ -0,0 +1,9 @@
+#!/usr/bin/env node
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+const agent_1 = __importDefault(require("./agent"));
+const manager = new agent_1.default();
+manager.start();

package/lib/index.d.ts

@@ -1,2 +1,2 @@
-#!/usr/bin/env node
-export {};
+export * from './agent';
+export { default } from './agent';

package/lib/index.js

@@ -1,9 +1,23 @@
-#!/usr/bin/env node
 "use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __exportStar = (this && this.__exportStar) || function(m, exports) {
+    for (var p in m) if (p !== "default" && !Object.prototype.hasOwnProperty.call(exports, p)) __createBinding(exports, m, p);
+};
 var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-const agent_1 = __importDefault(require("./agent"));
-const manager = new agent_1.default();
-manager.start();
+exports.default = void 0;
+__exportStar(require("./agent"), exports);
+var agent_1 = require("./agent");
+Object.defineProperty(exports, "default", { enumerable: true, get: function () { return __importDefault(agent_1).default; } });

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@saber2pr/ai-agent",
-  "version": "0.0.4",
+  "version": "0.0.6",
   "description": "AI Assistant CLI.",
   "author": "saber2pr",
   "license": "ISC",
@@ -8,7 +8,7 @@
     "lib"
   ],
   "bin": {
-    "sagent": "./lib/index.js"
+    "sagent": "./lib/cli.js"
   },
   "publishConfig": {
     "access": "public",
@@ -22,7 +22,8 @@
   },
   "dependencies": {
     "@modelcontextprotocol/sdk": "^1.25.3",
-    "@saber2pr/ts-context-mcp": "^0.0.5",
+    "@saber2pr/ts-context-mcp": "^0.0.6",
+    "js-tiktoken": "^1.0.21",
     "openai": "^6.16.0"
   },
   "devDependencies": {