From a68ff027b9041c5f365400a8f06667d3f8666c92 Mon Sep 17 00:00:00 2001 From: hex2077 Date: Thu, 7 Aug 2025 20:24:40 +0800 Subject: [PATCH] =?UTF-8?q?feat(oauth):=20=E6=B7=BB=E5=8A=A0OAuth=E4=BB=A4?= =?UTF-8?q?=E7=89=8C=E8=87=AA=E5=8A=A8=E5=88=B7=E6=96=B0=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为Gemini和Kiro服务添加令牌过期检测和自动刷新功能 在API服务器中添加定时任务检查并刷新临近过期的令牌 更新README文档说明新增的CRON配置参数 --- README-EN.md | 2 ++ README.md | 2 ++ src/adapter.js | 35 ++++++++++++++++++++++++- src/api-server.js | 54 ++++++++++++++++++++++++++++++++++----- src/claude/claude-kiro.js | 19 ++++++++++++++ src/gemini/gemini-core.js | 16 ++++++++++++ 6 files changed, 121 insertions(+), 7 deletions(-) diff --git a/README-EN.md b/README-EN.md index 3a19bdf..dd59f89 100644 --- a/README-EN.md +++ b/README-EN.md @@ -224,6 +224,8 @@ The following table details all supported parameters in `config.json`: | `PROMPT_LOG_BASE_NAME` | string | Base name for log files when `PROMPT_LOG_MODE` is `file`. | Defaults to `"prompt_log"` | | `REQUEST_MAX_RETRIES` | number | Maximum number of automatic retries for failed API requests. | Defaults to `3` | | `REQUEST_BASE_DELAY` | number | Base delay (milliseconds) between automatic retries. Delay increases with each retry. | Defaults to `1000` | +| `CRON_NEAR_MINUTES` | number | (Gemini-CLI Mode) Interval in minutes for the OAuth token refresh task. | Defaults to `15` | +| `CRON_REFRESH_TOKEN` | boolean | (Gemini-CLI Mode) Whether to enable automatic OAuth token refresh task. | Defaults to `true` | ### 3. Start the Service diff --git a/README.md b/README.md index 82f51f3..24c7e5b 100644 --- a/README.md +++ b/README.md @@ -224,6 +224,8 @@ claude-kiro-oauth。 | `PROMPT_LOG_BASE_NAME` | string | 当 `PROMPT_LOG_MODE` 为 `file` 时，生成的日志文件的基础名称。 | 默认为 `"prompt_log"` | | `REQUEST_MAX_RETRIES` | number | 当 API 请求失败时，自动重试的最大次数。 | 默认为 `3` | | `REQUEST_BASE_DELAY` | number | 自动重试之间的基础延迟时间（毫秒）。每次重试后延迟会增加。 | 默认为 `1000` | +| `CRON_NEAR_MINUTES` | number | (Gemini-CLI 模式) OAuth 令牌刷新任务计划的间隔时间（分钟）。 | 默认为 `15` | +| `CRON_REFRESH_TOKEN` | boolean | (Gemini-CLI 模式) 是否开启 OAuth 令牌自动刷新任务。 | 默认为 `true` | ### 3. 启动服务 diff --git a/src/adapter.js b/src/adapter.js index 06ce214..0445b73 100644 --- a/src/adapter.js +++ b/src/adapter.js @@ -40,6 +40,14 @@ export class ApiServiceAdapter { async listModels() { throw new Error("Method 'listModels()' must be implemented."); } + + /** + * 刷新认证令牌 + * @returns {Promise} + */ + async refreshToken() { + throw new Error("Method 'refreshToken()' must be implemented."); + } } // Gemini API 服务适配器 @@ -76,6 +84,14 @@ export class GeminiApiServiceAdapter extends ApiServiceAdapter { // Gemini Core API 的 listModels 已经返回符合 Gemini 格式的数据，所以不需要额外转换 return this.geminiApiService.listModels(); } + + async refreshToken() { + if(this.geminiApiService.isExpiryDateNear()===true){ + console.log(`[Gemini] Expiry date is near, refreshing token...`); + return this.geminiApiService.initializeAuth(true); + } + return Promise.resolve(); + } } // OpenAI API 服务适配器 @@ -102,6 +118,11 @@ export class OpenAIApiServiceAdapter extends ApiServiceAdapter { // The adapter now returns the native model list from the underlying service. return this.openAIApiService.listModels(); } + + async refreshToken() { + // OpenAI API keys are typically static and do not require refreshing. + return Promise.resolve(); + } } // Claude API 服务适配器 @@ -126,6 +147,10 @@ export class ClaudeApiServiceAdapter extends ApiServiceAdapter { // The adapter now returns the native model list from the underlying service. return this.claudeApiService.listModels(); } + + async refreshToken() { + return Promise.resolve(); + } } // Kiro API 服务适配器 @@ -153,10 +178,18 @@ export class KiroApiServiceAdapter extends ApiServiceAdapter { // Returns the native model list from the Kiro service return this.kiroApiService.listModels(); } + + async refreshToken() { + if(this.kiroApiService.isExpiryDateNear()===true){ + console.log(`[Kiro] Expiry date is near, refreshing token...`); + return this.kiroApiService.initializeAuth(true); + } + return Promise.resolve(); + } } // 用于存储服务适配器单例的映射 -const serviceInstances = {}; +export const serviceInstances = {}; // 服务适配器工厂 export function getServiceAdapter(config) { diff --git a/src/api-server.js b/src/api-server.js index 40d3b3e..c7e57bb 100644 --- a/src/api-server.js +++ b/src/api-server.js @@ -94,6 +94,10 @@ * --system-prompt-mode 系统提示模式 / System prompt mode: overwrite or append (default: overwrite) * --log-prompts 提示日志模式 / Prompt logging mode: console, file, or none (default: none) * --prompt-log-base-name 提示日志文件基础名称 / Base name for prompt log files (default: prompt_log) + * --request-max-retries API 请求失败时，自动重试的最大次数。 / Max retries for API requests on failure (default: 3) + * --request-base-delay 自动重试之间的基础延迟时间（毫秒）。每次重试后延迟会增加。 / Base delay in milliseconds between retries, increases with each retry (default: 1000) + * --cron-near-minutes OAuth 令牌刷新任务计划的间隔时间（分钟）。 / Interval for OAuth token refresh task in minutes (default: 15) + * --cron-refresh-token 是否开启 OAuth 令牌自动刷新任务 / Whether to enable automatic OAuth token refresh task (default: true) * */ @@ -105,7 +109,7 @@ import { promises as pfs } from 'fs'; import 'dotenv/config'; // Import dotenv and configure it import deepmerge from 'deepmerge'; -import { getServiceAdapter } from './adapter.js'; +import { getServiceAdapter, serviceInstances} from './adapter.js'; import { INPUT_SYSTEM_PROMPT_FILE, API_ACTIONS, @@ -155,7 +159,9 @@ async function initializeConfig(args = process.argv.slice(2), configFilePath = ' PROMPT_LOG_BASE_NAME: "prompt_log", PROMPT_LOG_MODE: "none", REQUEST_MAX_RETRIES: 3, - REQUEST_BASE_DELAY: 1000 + REQUEST_BASE_DELAY: 1000, + CRON_NEAR_MINUTES: 15, + CRON_REFRESH_TOKEN: true }; console.log('[Config] Using default configuration.'); } @@ -293,6 +299,20 @@ async function initializeConfig(args = process.argv.slice(2), configFilePath = ' } else { console.warn(`[Config Warning] --kiro-oauth-creds-file flag requires a value.`); } + } else if (args[i] === '--cron-near-minutes') { + if (i + 1 < args.length) { + currentConfig.CRON_NEAR_MINUTES = parseInt(args[i + 1], 10); + i++; + } else { + console.warn(`[Config Warning] --cron-near-minutes flag requires a value.`); + } + } else if (args[i] === '--cron-refresh-token') { + if (i + 1 < args.length) { + currentConfig.CRON_REFRESH_TOKEN = args[i + 1].toLowerCase() === 'true'; + i++; + } else { + console.warn(`[Config Warning] --cron-refresh-token flag requires a value.`); + } } } @@ -346,16 +366,17 @@ async function getSystemPromptFileContent(filePath) { } } -async function initApiService(config) { // Make getApiService exportable and accept config +async function initApiService(config) { // Initialize all known service adapters at startup for (const provider of Object.values(MODEL_PROVIDER)) { try { console.log(`[Initialization] Initializing service adapter for ${provider}...`); - getServiceAdapter({ ...config, MODEL_PROVIDER: provider }); + getServiceAdapter({ ...config, MODEL_PROVIDER: provider }); // This call populates serviceInstances } catch (error) { console.warn(`[Initialization Warning] Failed to initialize service adapter for ${provider}: ${error.message}`); } } + return serviceInstances; // Return the collection of initialized service instances } async function getApiService(config) { @@ -475,11 +496,25 @@ function createRequestHandler(config) { // --- Server Initialization --- async function startServer() { await initializeConfig(); // Initialize CONFIG globally - await initApiService(CONFIG); // Get service instance with the initialized CONFIG + const services = await initApiService(CONFIG); // Get service instance with the initialized CONFIG const requestHandlerInstance = createRequestHandler(CONFIG); // Create request handler with CONFIG and service - const server = http.createServer(requestHandlerInstance); + // 定义心跳和令牌刷新函数 + const heartbeatAndRefreshToken = async () => { + console.log(`[Heartbeat] Server is running. Current time: ${new Date().toLocaleString()}`); + // 循环遍历所有已初始化的服务适配器，并尝试刷新令牌 + for (const providerKey in services) { + const serviceAdapter = services[providerKey]; + try { + await serviceAdapter.refreshToken(); + console.log(`[Token Refresh] Refreshed token for ${providerKey}`); + } catch (error) { + console.error(`[Token Refresh Error] Failed to refresh token for ${providerKey}: ${error.message}`); + } + } + }; + const server = http.createServer(requestHandlerInstance); server.listen(CONFIG.SERVER_PORT, CONFIG.HOST, () => { console.log(`--- Unified API Server Configuration ---`); console.log(` Model Provider: ${CONFIG.MODEL_PROVIDER}`); @@ -508,6 +543,13 @@ async function startServer() { console.log(` • Gemini-compatible: /v1beta/models, /v1beta/models/{model}:generateContent`); console.log(` • Claude-compatible: /v1/messages`); console.log(` • Health check: /health`); + + if (CONFIG.CRON_REFRESH_TOKEN) { + console.log(` • Cron Near Minutes: ${CONFIG.CRON_NEAR_MINUTES}`); + console.log(` • Cron Refresh Token: ${CONFIG.CRON_REFRESH_TOKEN}`); + // 每 CRON_NEAR_MINUTES 分钟执行一次心跳日志和令牌刷新 + setInterval(heartbeatAndRefreshToken, CONFIG.CRON_NEAR_MINUTES * 60 * 1000); + } }); return server; // Return the server instance for testing purposes } diff --git a/src/claude/claude-kiro.js b/src/claude/claude-kiro.js index 3b9feae..885e34b 100644 --- a/src/claude/claude-kiro.js +++ b/src/claude/claude-kiro.js @@ -356,6 +356,7 @@ async initializeAuth(forceRefresh = false) { const filePath = path.join(dirPath, file); const credentials = await loadCredentialsFromFile(filePath); if (credentials) { + credentials.expiresAt = mergedCredentials.expiresAt; Object.assign(mergedCredentials, credentials); console.debug(`[Kiro Auth] Loaded credentials from ${file}`); } @@ -1062,4 +1063,22 @@ async initializeAuth(forceRefresh = false) { return { models: models }; } + + /** + * Checks if the given expiresAt timestamp is within 10 minutes from now. + * @returns {boolean} - True if expiresAt is less than 10 minutes from now, false otherwise. + */ + isExpiryDateNear() { + try { + const expirationTime = new Date(this.expiresAt); + const currentTime = new Date(); + const cronNearMinutesInMillis = (this.config.CRON_NEAR_MINUTES || 10) * 60 * 1000; + const thresholdTime = new Date(currentTime.getTime() + cronNearMinutesInMillis); + console.log(`[Kiro] Expiry date: ${expirationTime.getTime()}, Current time: ${currentTime.getTime()}, ${this.config.CRON_NEAR_MINUTES || 10} minutes from now: ${thresholdTime.getTime()}`); + return expirationTime.getTime() <= thresholdTime.getTime(); + } catch (error) { + console.error(`[Kiro] Error checking expiry date: ${this.expiresAt}, Error: ${error.message}`); + return false; // Treat as expired if parsing fails + } + } } diff --git a/src/gemini/gemini-core.js b/src/gemini/gemini-core.js index c4453f0..5302498 100644 --- a/src/gemini/gemini-core.js +++ b/src/gemini/gemini-core.js @@ -324,4 +324,20 @@ export class GeminiApiService { yield toGeminiApiResponse(chunk.response); } } + + /** + * Checks if the given expiry date is within the next 10 minutes from now. + * @returns {boolean} True if the expiry date is within the next 10 minutes, false otherwise. + */ + isExpiryDateNear() { + try { + const currentTime = Date.now(); + const cronNearMinutesInMillis = (this.config.CRON_NEAR_MINUTES || 10) * 60 * 1000; + console.log(`[Gemini] Expiry date: ${this.authClient.credentials.expiry_date}, Current time: ${currentTime}, ${this.config.CRON_NEAR_MINUTES || 10} minutes from now: ${currentTime + cronNearMinutesInMillis}`); + return this.authClient.credentials.expiry_date <= (currentTime + cronNearMinutesInMillis); + } catch (error) { + console.error(`[Gemini] Error checking expiry date: ${error.message}`); + return false; + } + } }