Merge pull request #95 from Sanyela/main

fix: 修复 Kiro API 工具调用、流式传输优化及号池错误处理

src/claude/claude-kiro.js

@@ -4,6 +4,8 @@ import { promises as fs } from 'fs';
 import * as path from 'path';
 import * as os from 'os';
 import * as crypto from 'crypto';
+import * as http from 'http';
+import * as https from 'https';
 import { getProviderModels } from '../provider-models.js';
 import { countTokens } from '@anthropic-ai/tokenizer';
 
@@ -281,8 +283,24 @@ export class KiroApiService {
         console.log('[Kiro] Initializing Kiro API Service...');
         await this.initializeAuth();
         const macSha256 = await getMacAddressSha256();
+        // 配置 HTTP/HTTPS agent 限制连接池大小，避免资源泄漏
+        const httpAgent = new http.Agent({
+            keepAlive: true,
+            maxSockets: 200,        // 每个主机最多 10 个连接
+            maxFreeSockets: 5,     // 最多保留 5 个空闲连接
+            timeout: 60000,        // 空闲连接 60 秒后关闭
+        });
+        const httpsAgent = new https.Agent({
+            keepAlive: true,
+            maxSockets: 100,
+            maxFreeSockets: 5,
+            timeout: 60000,
+        });
+        
         const axiosConfig = {
             timeout: KIRO_CONSTANTS.AXIOS_TIMEOUT,
+            httpAgent,
+            httpsAgent,
             headers: {
                 'Content-Type': KIRO_CONSTANTS.CONTENT_TYPE_JSON,
                 'x-amz-user-agent': `aws-sdk-js/1.0.7 KiroIDE-0.1.25-${macSha256}`,
@@ -1187,13 +1205,14 @@ async initializeAuth(forceRefresh = false) {
 
         const requestUrl = model.startsWith('amazonq') ? this.amazonQUrl : this.baseUrl;
 
+        let stream = null;
         try {
             const response = await this.axiosInstance.post(requestUrl, requestData, { 
                 headers,
                 responseType: 'stream'
             });
 
-            const stream = response.data;
+            stream = response.data;
             let buffer = '';
             let lastContentEvent = null;  // 用于检测连续重复的 content 事件
 
@@ -1224,6 +1243,11 @@ async initializeAuth(forceRefresh = false) {
                 }
             }
         } catch (error) {
+            // 确保出错时关闭流
+            if (stream && typeof stream.destroy === 'function') {
+                stream.destroy();
+            }
+            
             if (error.response?.status === 403 && !isRetry) {
                 console.log('[Kiro] Received 403 in stream. Attempting token refresh and retrying...');
                 await this.initializeAuth(true);
@@ -1241,6 +1265,11 @@ async initializeAuth(forceRefresh = false) {
 
             console.error('[Kiro] Stream API call failed:', error.message);
             throw error;
+        } finally {
+            // 确保流被关闭，释放资源
+            if (stream && typeof stream.destroy === 'function') {
+                stream.destroy();
+            }
         }
     }
 
@@ -1300,7 +1329,7 @@ async initializeAuth(forceRefresh = false) {
             for await (const event of this.streamApiReal('', finalModel, requestBody)) {
                 if (event.type === 'content' && event.content) {
                     totalContent += event.content;
-                    outputTokens += this.countTextTokens(event.content);
+                    // 不再每个 chunk 都计算 token，改为最后统一计算，避免阻塞事件循环
                     
                     yield {
                         type: "content_block_delta",
@@ -1412,12 +1441,16 @@ async initializeAuth(forceRefresh = false) {
                     };
                     
                     yield { type: "content_block_stop", index: blockIndex };
-                    
-                    outputTokens += this.countTextTokens(JSON.stringify(tc.input || {}));
                 }
             }
 
             // 6. 发送 message_delta 事件
+            // 在流结束后统一计算 output tokens，避免在流式循环中阻塞事件循环
+            outputTokens = this.countTextTokens(totalContent);
+            for (const tc of toolCalls) {
+                outputTokens += this.countTextTokens(JSON.stringify(tc.input || {}));
+            }
+            
             yield {
                 type: "message_delta",
                 delta: { stop_reason: toolCalls.length > 0 ? "tool_use" : "end_turn" },

src/common.js

@@ -402,9 +402,24 @@ export async function handleContentGenerationRequest(req, res, service, endpoint
 
     // 2.5. 如果使用了提供商池，根据模型重新选择提供商
     if (providerPoolManager && CONFIG.providerPools && CONFIG.providerPools[CONFIG.MODEL_PROVIDER]) {
-        const { getApiService } = await import('./service-manager.js');
-        service = await getApiService(CONFIG, model);
-        console.log(`[Content Generation] Re-selected service adapter based on model: ${model}`);
+        const { getApiService, NoAvailableProviderError } = await import('./service-manager.js');
+        try {
+            service = await getApiService(CONFIG, model);
+            console.log(`[Content Generation] Re-selected service adapter based on model: ${model}`);
+        } catch (error) {
+            if (error instanceof NoAvailableProviderError) {
+                // 号池无可用账号，返回 400 错误
+                res.writeHead(400, { 'Content-Type': 'application/json' });
+                res.end(JSON.stringify({
+                    error: {
+                        type: 'no_available_provider',
+                        message: error.message
+                    }
+                }));
+                return;
+            }
+            throw error;
+        }
     }
 
     // 3. Apply system prompt from file if configured.

src/service-manager.js

@@ -56,11 +56,31 @@ export async function initApiService(config) {
     return serviceInstances; // Return the collection of initialized service instances
 }
 
+/**
+ * Custom error class for no available provider in pool
+ */
+export class NoAvailableProviderError extends Error {
+    constructor(providerType, requestedModel = null) {
+        const message = requestedModel 
+            ? `号池无可用账号: ${providerType} (model: ${requestedModel})`
+            : `号池无可用账号: ${providerType}`;
+        super(message);
+        this.name = 'NoAvailableProviderError';
+        this.providerType = providerType;
+        this.requestedModel = requestedModel;
+        this.statusCode = 400;
+    }
+}
+
+// Provider types that should throw error when no available provider (instead of falling back to main config)
+const STRICT_POOL_PROVIDERS = ['claude-kiro-oauth'];
+
 /**
  * Get API service adapter, considering provider pools
  * @param {Object} config - The current request configuration
  * @param {string} [requestedModel] - Optional. The model name to filter providers by.
  * @returns {Promise<Object>} The API service adapter
+ * @throws {NoAvailableProviderError} If no healthy provider is available in the pool (only for strict pool providers)
  */
 export async function getApiService(config, requestedModel = null) {
     let serviceConfig = config;
@@ -74,7 +94,14 @@ export async function getApiService(config, requestedModel = null) {
             config.uuid = serviceConfig.uuid;
             console.log(`[API Service] Using pooled configuration for ${config.MODEL_PROVIDER}: ${serviceConfig.uuid}${requestedModel ? ` (model: ${requestedModel})` : ''}`);
         } else {
-            console.warn(`[API Service] No healthy provider found in pool for ${config.MODEL_PROVIDER}${requestedModel ? ` supporting model: ${requestedModel}` : ''}. Falling back to main config.`);
+            // 号池没有可用账号
+            // 只有 Kiro 类型的 provider 抛出错误，其他 provider 回退到主配置
+            if (STRICT_POOL_PROVIDERS.includes(config.MODEL_PROVIDER)) {
+                console.error(`[API Service] 号池无可用账号: ${config.MODEL_PROVIDER}${requestedModel ? ` (model: ${requestedModel})` : ''}`);
+                throw new NoAvailableProviderError(config.MODEL_PROVIDER, requestedModel);
+            } else {
+                console.warn(`[API Service] No healthy provider found in pool for ${config.MODEL_PROVIDER}${requestedModel ? ` supporting model: ${requestedModel}` : ''}. Falling back to main config.`);
+            }
         }
     }
     return getServiceAdapter(serviceConfig);