From 10c69a586bc20d83feec0372ad717c1100a5de8e Mon Sep 17 00:00:00 2001 From: musistudio Date: Wed, 31 Dec 2025 22:44:16 +0800 Subject: [PATCH] add token speed block --- packages/cli/package.json | 14 +- packages/cli/src/utils/statusline.ts | 500 ++++++++++++----- packages/core/src/plugins/index.ts | 2 +- packages/core/src/plugins/output/index.ts | 15 + .../core/src/plugins/output/output-manager.ts | 4 + .../src/plugins/output/temp-file-handler.ts | 140 +++++ packages/core/src/plugins/output/types.ts | 29 +- packages/core/src/plugins/token-speed.ts | 501 ++++++++++-------- packages/core/src/server.ts | 2 +- packages/server/src/index.ts | 6 + packages/server/src/middleware/auth.ts | 3 +- packages/server/src/types.d.ts | 42 ++ .../src/components/StatusLineConfigDialog.tsx | 9 + scripts/build-cli.js | 2 +- 14 files changed, 900 insertions(+), 369 deletions(-) create mode 100644 packages/core/src/plugins/output/temp-file-handler.ts diff --git a/packages/cli/package.json b/packages/cli/package.json index a905fee..ca92912 100644 --- a/packages/cli/package.json +++ b/packages/cli/package.json @@ -17,20 +17,18 @@ ], "author": "musistudio", "license": "MIT", - "dependencies": { + "devDependencies": { "@CCR/server": "workspace:*", "@CCR/shared": "workspace:*", "@inquirer/prompts": "^5.0.0", - "adm-zip": "^0.5.16", - "archiver": "^7.0.1", - "find-process": "^2.0.0", - "minimist": "^1.2.8", - "openurl": "^1.1.1" - }, - "devDependencies": { "@types/archiver": "^7.0.0", "@types/node": "^24.0.15", + "adm-zip": "^0.5.16", + "archiver": "^7.0.1", "esbuild": "^0.25.1", + "find-process": "^2.0.0", + "minimist": "^1.2.8", + "openurl": "^1.1.1", "ts-node": "^10.9.2", "typescript": "^5.8.2" } diff --git a/packages/cli/src/utils/statusline.ts b/packages/cli/src/utils/statusline.ts index 8bbe89b..2567068 100644 --- a/packages/cli/src/utils/statusline.ts +++ b/packages/cli/src/utils/statusline.ts @@ -1,7 +1,8 @@ import fs from "node:fs/promises"; import path from "node:path"; import { execSync } from "child_process"; -import { CONFIG_FILE } from "@CCR/shared"; +import { tmpdir } from "node:os"; +import { CONFIG_FILE, HOME_DIR } from "@CCR/shared"; import JSON5 from "json5"; export interface StatusLineModuleConfig { @@ -10,7 +11,7 @@ export interface StatusLineModuleConfig { text: string; color?: string; background?: string; - scriptPath?: string; // 用于script类型的模块，指定要执行的Node.js脚本文件路径 + scriptPath?: string; // Path to the Node.js script file to execute for script-type modules } export interface StatusLineThemeConfig { @@ -30,6 +31,28 @@ export interface StatusLineInput { current_dir: string; project_dir: string; }; + version?: string; + output_style?: { + name: string; + }; + cost?: { + total_cost_usd: number; + total_duration_ms: number; + total_api_duration_ms: number; + total_lines_added: number; + total_lines_removed: number; + }; + context_window?: { + total_input_tokens: number; + total_output_tokens: number; + context_window_size: number; + current_usage: { + input_tokens: number; + output_tokens: number; + cache_creation_input_tokens: number; + cache_read_input_tokens: number; + } | null; + }; } export interface AssistantMessage { @@ -43,12 +66,12 @@ export interface AssistantMessage { }; } -// ANSIColor代码 +// ANSI Color codes const COLORS: Record = { reset: "\x1b[0m", bold: "\x1b[1m", dim: "\x1b[2m", - // 标准颜色 + // Standard colors black: "\x1b[30m", red: "\x1b[31m", green: "\x1b[32m", @@ -57,7 +80,7 @@ const COLORS: Record = { magenta: "\x1b[35m", cyan: "\x1b[36m", white: "\x1b[37m", - // 亮色 + // Bright colors bright_black: "\x1b[90m", bright_red: "\x1b[91m", bright_green: "\x1b[92m", @@ -66,7 +89,7 @@ const COLORS: Record = { bright_magenta: "\x1b[95m", bright_cyan: "\x1b[96m", bright_white: "\x1b[97m", - // 背景颜色 + // Background colors bg_black: "\x1b[40m", bg_red: "\x1b[41m", bg_green: "\x1b[42m", @@ -75,7 +98,7 @@ const COLORS: Record = { bg_magenta: "\x1b[45m", bg_cyan: "\x1b[46m", bg_white: "\x1b[47m", - // 亮背景色 + // Bright background colors bg_bright_black: "\x1b[100m", bg_bright_red: "\x1b[101m", bg_bright_green: "\x1b[102m", @@ -86,16 +109,16 @@ const COLORS: Record = { bg_bright_white: "\x1b[107m", }; -// 使用TrueColor(24位色)支持十六进制颜色 +// Use TrueColor (24-bit color) to support hexadecimal colors const TRUE_COLOR_PREFIX = "\x1b[38;2;"; const TRUE_COLOR_BG_PREFIX = "\x1b[48;2;"; -// 将十六进制颜色转为RGB格式 +// Convert hexadecimal color to RGB format function hexToRgb(hex: string): { r: number; g: number; b: number } | null { - // 移除#和空格 + // Remove # and spaces hex = hex.replace(/^#/, '').trim(); - // 处理简写形式 (#RGB -> #RRGGBB) + // Handle shorthand form (#RGB -> #RRGGBB) if (hex.length === 3) { hex = hex[0] + hex[0] + hex[1] + hex[1] + hex[2] + hex[2]; } @@ -108,7 +131,7 @@ function hexToRgb(hex: string): { r: number; g: number; b: number } | null { const g = parseInt(hex.substring(2, 4), 16); const b = parseInt(hex.substring(4, 6), 16); - // 验证RGB值是否有效 + // Validate RGB values if (isNaN(r) || isNaN(g) || isNaN(b) || r < 0 || r > 255 || g < 0 || g > 255 || b < 0 || b > 255) { return null; } @@ -116,9 +139,9 @@ function hexToRgb(hex: string): { r: number; g: number; b: number } | null { return { r, g, b }; } -// 获取颜色代码 +// Get color code function getColorCode(colorName: string): string { - // 检查是否是十六进制颜色 + // Check if it's a hexadecimal color if (colorName.startsWith('#') || /^[0-9a-fA-F]{6}$/.test(colorName) || /^[0-9a-fA-F]{3}$/.test(colorName)) { const rgb = hexToRgb(colorName); if (rgb) { @@ -126,66 +149,66 @@ function getColorCode(colorName: string): string { } } - // 默认返回空字符串 + // Default to empty string return ""; } -// 变量替换函数，支持{{var}}格式的变量替换 +// Variable replacement function, supports {{var}} format variable replacement function replaceVariables(text: string, variables: Record): string { return text.replace(/\{\{(\w+)\}\}/g, (_match, varName) => { return variables[varName] || ""; }); } -// 执行脚本并获取输出 +// Execute script and get output async function executeScript(scriptPath: string, variables: Record): Promise { try { - // 检查文件是否存在 + // Check if file exists await fs.access(scriptPath); - // 使用require动态加载脚本模块 + // Use require to dynamically load script module const scriptModule = require(scriptPath); - // 如果导出的是函数，则调用它并传入变量 + // If export is a function, call it with variables if (typeof scriptModule === 'function') { const result = scriptModule(variables); - // 如果返回的是Promise，则等待它完成 + // If returns a Promise, wait for it to complete if (result instanceof Promise) { return await result; } return result; } - // 如果导出的是default函数，则调用它 + // If export is a default function, call it if (scriptModule.default && typeof scriptModule.default === 'function') { const result = scriptModule.default(variables); - // 如果返回的是Promise，则等待它完成 + // If returns a Promise, wait for it to complete if (result instanceof Promise) { return await result; } return result; } - // 如果导出的是字符串，则直接返回 + // If export is a string, return directly if (typeof scriptModule === 'string') { return scriptModule; } - // 如果导出的是default字符串，则返回它 + // If export is a default string, return it if (scriptModule.default && typeof scriptModule.default === 'string') { return scriptModule.default; } - // 默认情况下返回空字符串 + // Default to empty string return ""; } catch (error) { - console.error(`执行脚本 ${scriptPath} 时出错:`, error); + console.error(`Error executing script ${scriptPath}:`, error); return ""; } } -// 默认主题配置 - 使用Nerd Fonts图标和美观配色 +// Default theme configuration - using Nerd Fonts icons and beautiful color scheme const DEFAULT_THEME: StatusLineThemeConfig = { modules: [ { @@ -208,20 +231,20 @@ const DEFAULT_THEME: StatusLineThemeConfig = { }, { type: "usage", - icon: "↑", // 上箭头 + icon: "↑", // Up arrow text: "{{inputTokens}}", color: "bright_green" }, { type: "usage", - icon: "↓", // 下箭头 + icon: "↓", // Down arrow text: "{{outputTokens}}", color: "bright_yellow" } ] }; -// Powerline风格主题配置 +// Powerline style theme configuration const POWERLINE_THEME: StatusLineThemeConfig = { modules: [ { @@ -247,14 +270,14 @@ const POWERLINE_THEME: StatusLineThemeConfig = { }, { type: "usage", - icon: "↑", // 上箭头 + icon: "↑", // Up arrow text: "{{inputTokens}}", color: "white", background: "bg_bright_green" }, { type: "usage", - icon: "↓", // 下箭头 + icon: "↓", // Down arrow text: "{{outputTokens}}", color: "white", background: "bg_bright_yellow" @@ -262,7 +285,7 @@ const POWERLINE_THEME: StatusLineThemeConfig = { ] }; -// 简单文本主题配置 - 用于图标无法显示时的fallback +// Simple text theme configuration - fallback for when icons cannot be displayed const SIMPLE_THEME: StatusLineThemeConfig = { modules: [ { @@ -298,7 +321,61 @@ const SIMPLE_THEME: StatusLineThemeConfig = { ] }; -// 格式化usage信息，如果大于1000则使用k单位 +// Full theme configuration - showcasing all available modules +const FULL_THEME: StatusLineThemeConfig = { + modules: [ + { + type: "workDir", + icon: "󰉋", + text: "{{workDirName}}", + color: "bright_blue" + }, + { + type: "gitBranch", + icon: "", + text: "{{gitBranch}}", + color: "bright_magenta" + }, + { + type: "model", + icon: "󰚩", + text: "{{model}}", + color: "bright_cyan" + }, + { + type: "context", + icon: "🪟", + text: "{{contextPercent}}% / {{contextWindowSize}}", + color: "bright_green" + }, + { + type: "speed", + icon: "⚡", + text: "{{tokenSpeed}} t/s {{isStreaming}}", + color: "bright_yellow" + }, + { + type: "cost", + icon: "💰", + text: "{{cost}}", + color: "bright_magenta" + }, + { + type: "duration", + icon: "⏱️", + text: "{{duration}}", + color: "bright_white" + }, + { + type: "lines", + icon: "📝", + text: "+{{linesAdded}}/-{{linesRemoved}}", + color: "bright_cyan" + } + ] +}; + +// Format usage information, use k unit if greater than 1000 function formatUsage(input_tokens: number, output_tokens: number): string { if (input_tokens > 1000 || output_tokens > 1000) { const inputFormatted = input_tokens > 1000 ? `${(input_tokens / 1000).toFixed(1)}k` : `${input_tokens}`; @@ -308,13 +385,126 @@ function formatUsage(input_tokens: number, output_tokens: number): string { return `${input_tokens} ${output_tokens}`; } -// 读取用户主目录的主题配置 +// Calculate context window usage percentage +function calculateContextPercent(context_window: StatusLineInput['context_window']): number { + if (!context_window || !context_window.current_usage) { + return 0; + } + const { current_usage, context_window_size } = context_window; + const currentTokens = current_usage.input_tokens + + current_usage.cache_creation_input_tokens + + current_usage.cache_read_input_tokens; + return Math.round((currentTokens / context_window_size) * 100); +} + +// Format cost display +function formatCost(cost_usd: number): string { + if (cost_usd < 0.01) { + return `${(cost_usd * 100).toFixed(2)}¢`; + } + return `$${cost_usd.toFixed(2)}`; +} + +// Format duration +function formatDuration(ms: number): string { + if (Number.isNaN(ms)) { + return '' + } + if (ms < 1000) { + return `${ms}ms`; + } else if (ms < 60000) { + return `${(ms / 1000).toFixed(1)}s`; + } else { + const minutes = Math.floor(ms / 60000); + const seconds = ((ms % 60000) / 1000).toFixed(0); + if (Number.isNaN(minutes) || Number.isNaN(seconds)) { + return '' + } + return `${minutes}m${seconds}s`; + } +} + +// Read token-speed statistics from temp file +async function getTokenSpeedStats(sessionId: string): Promise<{ + tokensPerSecond: number; + timeToFirstToken?: number; +} | null> { + try { + // Use system temp directory + const tempDir = path.join(tmpdir(), 'claude-code-router'); + + // Check if temp directory exists + try { + await fs.access(tempDir); + } catch { + return null; + } + + // List all files in temp directory + const files = await fs.readdir(tempDir); + + // Find files matching pattern: session-{sessionId}-{timestamp}.json + const pattern = new RegExp(`^session-${sessionId}-(\\d+)\\.json$`); + const matchingFiles = files + .map(file => { + const match = file.match(pattern); + if (!match) return null; + return { + file, + timestamp: parseInt(match[1]) + }; + }) + .filter(Boolean) as Array<{ file: string; timestamp: number }>; + + if (matchingFiles.length === 0) { + return null; + } + + // Sort by timestamp descending and get the most recent file + matchingFiles.sort((a, b) => b.timestamp - a.timestamp); + const latestFile = matchingFiles[0]; + const statsFilePath = path.join(tempDir, latestFile.file); + + // Read stats file + const content = await fs.readFile(statsFilePath, 'utf-8'); + const data = JSON.parse(content); + + // Check if data has tokensPerSecond + if (data.tokensPerSecond !== undefined && data.tokensPerSecond > 0) { + // Check if timestamp is within last 3 seconds + const now = Date.now(); + const timestamp = data.timestamp || 0; + const ageInSeconds = (now - timestamp) / 1000; + + // If data is older than 3 seconds, return 0 speed + if (ageInSeconds > 3) { + return { + tokensPerSecond: 0, + timeToFirstToken: data.timeToFirstToken + }; + } + + const result = { + tokensPerSecond: parseInt(data.tokensPerSecond), + timeToFirstToken: data.timeToFirstToken + }; + return result; + } + + return null; + } catch (error) { + // Silently fail on error + return null; + } +} + +// Read theme configuration from user home directory async function getProjectThemeConfig(): Promise<{ theme: StatusLineThemeConfig | null, style: string }> { try { - // 只使用主目录的固定配置文件 + // Only use fixed configuration file in home directory const configPath = CONFIG_FILE; - // 检查配置文件是否存在 + // Check if configuration file exists try { await fs.access(configPath); } catch { @@ -324,52 +514,52 @@ async function getProjectThemeConfig(): Promise<{ theme: StatusLineThemeConfig | const configContent = await fs.readFile(configPath, "utf-8"); const config = JSON5.parse(configContent); - // 检查是否有StatusLine配置 + // Check if there's StatusLine configuration if (config.StatusLine) { - // 获取当前使用的风格，默认为default + // Get current style, default to 'default' const currentStyle = config.StatusLine.currentStyle || 'default'; - // 检查是否有对应风格的配置 + // Check if there's configuration for the corresponding style if (config.StatusLine[currentStyle] && config.StatusLine[currentStyle].modules) { return { theme: config.StatusLine[currentStyle], style: currentStyle }; } } } catch (error) { - // 如果读取失败，返回null + // Return null if reading fails // console.error("Failed to read theme config:", error); } return { theme: null, style: 'default' }; } -// 检查是否应该使用简单主题（fallback方案） -// 当环境变量 USE_SIMPLE_ICONS 被设置时，或者当检测到可能不支持Nerd Fonts的终端时 +// Check if simple theme should be used (fallback scheme) +// When environment variable USE_SIMPLE_ICONS is set, or when a terminal that might not support Nerd Fonts is detected function shouldUseSimpleTheme(): boolean { - // 检查环境变量 + // Check environment variable if (process.env.USE_SIMPLE_ICONS === 'true') { return true; } - // 检查终端类型（一些常见的不支持复杂图标的终端） + // Check terminal type (some common terminals that don't support complex icons) const term = process.env.TERM || ''; const unsupportedTerms = ['dumb', 'unknown']; if (unsupportedTerms.includes(term)) { return true; } - // 默认情况下，假设终端支持Nerd Fonts + // By default, assume terminal supports Nerd Fonts return false; } -// 检查Nerd Fonts图标是否能正确显示 -// 通过检查终端字体信息或使用试探性方法 +// Check if Nerd Fonts icons can be displayed correctly +// By checking terminal font information or using heuristic methods function canDisplayNerdFonts(): boolean { - // 如果环境变量明确指定使用简单图标，则不能显示Nerd Fonts + // If environment variable explicitly specifies simple icons, Nerd Fonts cannot be displayed if (process.env.USE_SIMPLE_ICONS === 'true') { return false; } - // 检查一些常见的支持Nerd Fonts的终端环境变量 + // Check some common terminal environment variables that support Nerd Fonts const fontEnvVars = ['NERD_FONT', 'NERDFONT', 'FONT']; for (const envVar of fontEnvVars) { const value = process.env[envVar]; @@ -378,36 +568,36 @@ function canDisplayNerdFonts(): boolean { } } - // 检查终端类型 + // Check terminal type const termProgram = process.env.TERM_PROGRAM || ''; const supportedTerminals = ['iTerm.app', 'vscode', 'Hyper', 'kitty', 'alacritty']; if (supportedTerminals.includes(termProgram)) { return true; } - // 检查COLORTERM环境变量 + // Check COLORTERM environment variable const colorTerm = process.env.COLORTERM || ''; if (colorTerm.includes('truecolor') || colorTerm.includes('24bit')) { return true; } - // 默认情况下，假设可以显示Nerd Fonts（但允许用户通过环境变量覆盖） + // By default, assume Nerd Fonts can be displayed (but allow users to override via environment variables) return process.env.USE_SIMPLE_ICONS !== 'true'; } -// 检查特定Unicode字符是否能正确显示 -// 这是一个简单的试探性检查 +// Check if specific Unicode characters can be displayed correctly +// This is a simple heuristic check function canDisplayUnicodeCharacter(char: string): boolean { - // 对于Nerd Fonts图标，我们假设支持UTF-8的终端可以显示 - // 但实际上很难准确检测，所以我们依赖环境变量和终端类型检测 + // For Nerd Font icons, we assume UTF-8 terminals can display them + // But accurate detection is difficult, so we rely on environment variables and terminal type detection try { - // 检查终端是否支持UTF-8 + // Check if terminal supports UTF-8 const lang = process.env.LANG || process.env.LC_ALL || process.env.LC_CTYPE || ''; if (lang.includes('UTF-8') || lang.includes('utf8') || lang.includes('UTF8')) { return true; } - // 检查LC_*环境变量 + // Check LC_* environment variables const lcVars = ['LC_ALL', 'LC_CTYPE', 'LANG']; for (const lcVar of lcVars) { const value = process.env[lcVar]; @@ -416,35 +606,35 @@ function canDisplayUnicodeCharacter(char: string): boolean { } } } catch (e) { - // 如果检查失败，默认返回true + // If check fails, default to true return true; } - // 默认情况下，假设可以显示 + // By default, assume it can be displayed return true; } export async function parseStatusLineData(input: StatusLineInput): Promise { try { - // 检查是否应该使用简单主题 + // Check if simple theme should be used const useSimpleTheme = shouldUseSimpleTheme(); - // 检查是否可以显示Nerd Fonts图标 + // Check if Nerd Fonts icons can be displayed const canDisplayNerd = canDisplayNerdFonts(); - // 确定使用的主题：如果用户强制使用简单主题或无法显示Nerd Fonts，则使用简单主题 + // Determine which theme to use: use simple theme if user forces it or Nerd Fonts cannot be displayed const effectiveTheme = useSimpleTheme || !canDisplayNerd ? SIMPLE_THEME : DEFAULT_THEME; - // 获取主目录的主题配置，如果没有则使用确定的默认配置 + // Get theme configuration from home directory, or use the determined default configuration const { theme: projectTheme, style: currentStyle } = await getProjectThemeConfig(); const theme = projectTheme || effectiveTheme; - // 获取当前工作目录和Git分支 + // Get current working directory and Git branch const workDir = input.workspace.current_dir; let gitBranch = ""; try { - // 尝试获取Git分支名 + // Try to get Git branch name gitBranch = execSync("git branch --show-current", { cwd: workDir, stdio: ["pipe", "pipe", "ignore"], @@ -452,14 +642,14 @@ export async function parseStatusLineData(input: StatusLineInput): Promise 0 + ? tokenSpeedData.tokensPerSecond.toString() + : ''; + + // Check if streaming (has active token speed) + const isStreaming = tokenSpeedData !== null && tokenSpeedData.tokensPerSecond > 0; + + const streamingIndicator = isStreaming ? '[Streaming]' : '' + + // Format time to first token + let formattedTimeToFirstToken = ''; + if (tokenSpeedData?.timeToFirstToken !== undefined) { + formattedTimeToFirstToken = formatDuration(tokenSpeedData.timeToFirstToken); + } + + // Process context window data + const contextPercent = input.context_window ? calculateContextPercent(input.context_window) : 0; + const totalInputTokens = input.context_window?.total_input_tokens || 0; + const totalOutputTokens = input.context_window?.total_output_tokens || 0; + const contextWindowSize = input.context_window?.context_window_size || 0; + + // Process cost data + const totalCost = input.cost?.total_cost_usd || 0; + const formattedCost = totalCost > 0 ? formatCost(totalCost) : ''; + const totalDuration = input.cost?.total_duration_ms || 0; + const formattedDuration = totalDuration > 0 ? formatDuration(totalDuration) : ''; + const linesAdded = input.cost?.total_lines_added || 0; + const linesRemoved = input.cost?.total_lines_removed || 0; + + // Define variable replacement mapping + const variables: Record = { workDirName, gitBranch, model, inputTokens: formattedInputTokens, - outputTokens: formattedOutputTokens + outputTokens: formattedOutputTokens, + tokenSpeed: formattedTokenSpeed || '0', + isStreaming: isStreaming ? 'streaming' : '', + timeToFirstToken: formattedTimeToFirstToken, + contextPercent: contextPercent.toString(), + streamingIndicator, + contextWindowSize: contextWindowSize > 1000 ? `${(contextWindowSize / 1000).toFixed(0)}k` : contextWindowSize.toString(), + totalInputTokens: totalInputTokens > 1000 ? `${(totalInputTokens / 1000).toFixed(1)}k` : totalInputTokens.toString(), + totalOutputTokens: totalOutputTokens > 1000 ? `${(totalOutputTokens / 1000).toFixed(1)}k` : totalOutputTokens.toString(), + cost: formattedCost || '', + duration: formattedDuration || '', + linesAdded: linesAdded.toString(), + linesRemoved: linesRemoved.toString(), + netLines: (linesAdded - linesRemoved).toString(), + version: input.version || '', + sessionId: input.session_id.substring(0, 8) }; - // 确定使用的风格 + // Determine the style to use const isPowerline = currentStyle === 'powerline'; - // 根据风格渲染状态行 + // Render status line based on style if (isPowerline) { return await renderPowerlineStyle(theme, variables); } else { return await renderDefaultStyle(theme, variables); } } catch (error) { - // 发生错误时返回空字符串 + // Return empty string on error return ""; } } -// 读取用户主目录的主题配置（指定风格） +// Read theme configuration from user home directory (specified style) async function getProjectThemeConfigForStyle(style: string): Promise { try { - // 只使用主目录的固定配置文件 + // Only use fixed configuration file in home directory const configPath = CONFIG_FILE; - // 检查配置文件是否存在 + // Check if configuration file exists try { await fs.access(configPath); } catch { @@ -564,19 +800,19 @@ async function getProjectThemeConfigForStyle(style: string): Promise @@ -584,14 +820,15 @@ async function renderDefaultStyle( const modules = theme.modules || DEFAULT_THEME.modules; const parts: string[] = []; - // 遍历模块数组，渲染每个模块 - for (let i = 0; i < Math.min(modules.length, 5); i++) { + // Iterate through module array, rendering each module (maximum 10) + for (let i = 0; i < modules.length; i++) { const module = modules[i]; + const color = module.color ? getColorCode(module.color) : ""; const background = module.background ? getColorCode(module.background) : ""; const icon = module.icon || ""; - // 如果是script类型，执行脚本获取文本 + // If script type, execute script to get text let text = ""; if (module.type === "script" && module.scriptPath) { text = await executeScript(module.scriptPath, variables); @@ -599,35 +836,35 @@ async function renderDefaultStyle( text = replaceVariables(module.text, variables); } - // 构建显示文本 + // Build display text let displayText = ""; if (icon) { displayText += `${icon} `; } displayText += text; - // 如果displayText为空，或者只有图标没有实际文本，则跳过该模块 + // Skip module if displayText is empty or only has icon without actual text if (!displayText || !text) { continue; } - // 构建模块字符串 + // Build module string let part = `${background}${color}`; part += `${displayText}${COLORS.reset}`; parts.push(part); } - // 使用空格连接所有部分 + // Join all parts with spaces return parts.join(" "); } -// Powerline符号 +// Powerline symbols const SEP_RIGHT = "\uE0B0"; //  -// 颜色编号（256色表） +// Color numbers (256-color table) const COLOR_MAP: Record = { - // 基础颜色映射到256色 + // Basic colors mapped to 256 colors black: 0, red: 1, green: 2, @@ -644,7 +881,7 @@ const COLOR_MAP: Record = { bright_magenta: 13, bright_cyan: 14, bright_white: 15, - // 亮背景色映射 + // Bright background color mapping bg_black: 0, bg_red: 1, bg_green: 2, @@ -661,25 +898,25 @@ const COLOR_MAP: Record = { bg_bright_magenta: 13, bg_bright_cyan: 14, bg_bright_white: 15, - // 自定义颜色映射 + // Custom color mapping bg_bright_orange: 202, bg_bright_purple: 129, }; -// 获取TrueColor的RGB值 +// Get TrueColor RGB value function getTrueColorRgb(colorName: string): { r: number; g: number; b: number } | null { - // 如果是预定义颜色，返回对应RGB + // If predefined color, return corresponding RGB if (COLOR_MAP[colorName] !== undefined) { const color256 = COLOR_MAP[colorName]; return color256ToRgb(color256); } - // 处理十六进制颜色 + // Handle hexadecimal color if (colorName.startsWith('#') || /^[0-9a-fA-F]{6}$/.test(colorName) || /^[0-9a-fA-F]{3}$/.test(colorName)) { return hexToRgb(colorName); } - // 处理背景色十六进制 + // Handle background color hexadecimal if (colorName.startsWith('bg_#')) { return hexToRgb(colorName.substring(3)); } @@ -687,13 +924,13 @@ function getTrueColorRgb(colorName: string): { r: number; g: number; b: number } return null; } -// 将256色表索引转换为RGB值 +// Convert 256-color table index to RGB value function color256ToRgb(index: number): { r: number; g: number; b: number } | null { if (index < 0 || index > 255) return null; - // ANSI 256色表转换 + // ANSI 256-color table conversion if (index < 16) { - // 基本颜色 + // Basic colors const basicColors = [ [0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128], [128, 0, 128], [0, 128, 128], [192, 192, 192], @@ -702,7 +939,7 @@ function color256ToRgb(index: number): { r: number; g: number; b: number } | nul ]; return { r: basicColors[index][0], g: basicColors[index][1], b: basicColors[index][2] }; } else if (index < 232) { - // 216色：6×6×6的颜色立方体 + // 216 colors: 6×6×6 color cube const i = index - 16; const r = Math.floor(i / 36); const g = Math.floor((i % 36) / 6); @@ -710,17 +947,17 @@ function color256ToRgb(index: number): { r: number; g: number; b: number } | nul const rgb = [0, 95, 135, 175, 215, 255]; return { r: rgb[r], g: rgb[g], b: rgb[b] }; } else { - // 灰度色 + // Grayscale colors const gray = 8 + (index - 232) * 10; return { r: gray, g: gray, b: gray }; } } -// 生成一个无缝拼接的段：文本在 bgN 上显示，分隔符从 bgN 过渡到 nextBgN +// Generate a seamless segment: text displayed on bgN, separator transitions from bgN to nextBgN function segment(text: string, textFg: string, bgColor: string, nextBgColor: string | null): string { const bgRgb = getTrueColorRgb(bgColor); if (!bgRgb) { - // 如果无法获取RGB，使用默认蓝色背景 + // If unable to get RGB, use default blue background const defaultBlueRgb = { r: 33, g: 150, b: 243 }; const curBg = `\x1b[48;2;${defaultBlueRgb.r};${defaultBlueRgb.g};${defaultBlueRgb.b}m`; const fgColor = `\x1b[38;2;255;255;255m`; @@ -730,8 +967,8 @@ function segment(text: string, textFg: string, bgColor: string, nextBgColor: str const curBg = `\x1b[48;2;${bgRgb.r};${bgRgb.g};${bgRgb.b}m`; - // 获取前景色RGB - let fgRgb = { r: 255, g: 255, b: 255 }; // 默认前景色为白色 + // Get foreground color RGB + let fgRgb = { r: 255, g: 255, b: 255 }; // Default foreground color is white const textFgRgb = getTrueColorRgb(textFg); if (textFgRgb) { fgRgb = textFgRgb; @@ -743,15 +980,15 @@ function segment(text: string, textFg: string, bgColor: string, nextBgColor: str if (nextBgColor != null) { const nextBgRgb = getTrueColorRgb(nextBgColor); if (nextBgRgb) { - // 分隔符：前景色是当前段的背景色，背景色是下一段的背景色 + // Separator: foreground color is current segment's background color, background color is next segment's background color const sepCurFg = `\x1b[38;2;${bgRgb.r};${bgRgb.g};${bgRgb.b}m`; const sepNextBg = `\x1b[48;2;${nextBgRgb.r};${nextBgRgb.g};${nextBgRgb.b}m`; const sep = `${sepCurFg}${sepNextBg}${SEP_RIGHT}\x1b[0m`; return body + sep; } - // 如果没有下一个背景色，假设终端背景为黑色并渲染黑色箭头 + // If no next background color, assume terminal background is black and render black arrow const sepCurFg = `\x1b[38;2;${bgRgb.r};${bgRgb.g};${bgRgb.b}m`; - const sepNextBg = `\x1b[48;2;0;0;0m`; // 黑色背景 + const sepNextBg = `\x1b[48;2;0;0;0m`; // Black background const sep = `${sepCurFg}${sepNextBg}${SEP_RIGHT}\x1b[0m`; return body + sep; } @@ -759,7 +996,7 @@ function segment(text: string, textFg: string, bgColor: string, nextBgColor: str return body; } -// 渲染Powerline风格的状态行 +// Render Powerline style status line async function renderPowerlineStyle( theme: StatusLineThemeConfig, variables: Record @@ -767,44 +1004,47 @@ async function renderPowerlineStyle( const modules = theme.modules || POWERLINE_THEME.modules; const segments: string[] = []; - // 遍历模块数组，渲染每个模块 - for (let i = 0; i < Math.min(modules.length, 5); i++) { + // Iterate through module array, rendering each module (maximum 10) + for (let i = 0; i < Math.min(modules.length, 10); i++) { const module = modules[i]; const color = module.color || "white"; const backgroundName = module.background || ""; const icon = module.icon || ""; - // 如果是script类型，执行脚本获取文本 + // If script type, execute script to get text let text = ""; if (module.type === "script" && module.scriptPath) { text = await executeScript(module.scriptPath, variables); + } else if (module.type === "speed") { + // speed module: use tokenSpeed variable + text = replaceVariables(module.text, variables); } else { text = replaceVariables(module.text, variables); } - // 构建显示文本 + // Build display text let displayText = ""; if (icon) { displayText += `${icon} `; } displayText += text; - // 如果displayText为空，或者只有图标没有实际文本，则跳过该模块 + // Skip module if displayText is empty or only has icon without actual text if (!displayText || !text) { continue; } - // 获取下一个模块的背景色（用于分隔符） + // Get next module's background color (for separator) let nextBackground: string | null = null; if (i < modules.length - 1) { const nextModule = modules[i + 1]; nextBackground = nextModule.background || null; } - // 使用模块定义的背景色，或者为Powerline风格提供默认背景色 + // Use module-defined background color, or provide default background color for Powerline style const actualBackground = backgroundName || "bg_bright_blue"; - // 生成段，支持十六进制颜色 + // Generate segment, supports hexadecimal colors const segmentStr = segment(displayText, color, actualBackground, nextBackground); segments.push(segmentStr); } diff --git a/packages/core/src/plugins/index.ts b/packages/core/src/plugins/index.ts index 0385647..ee8b868 100644 --- a/packages/core/src/plugins/index.ts +++ b/packages/core/src/plugins/index.ts @@ -1,4 +1,4 @@ export * from './types'; export { pluginManager } from './plugin-manager'; -export { tokenSpeedPlugin } from './token-speed'; +export { tokenSpeedPlugin, getTokenSpeedStats, getGlobalTokenSpeedStats } from './token-speed'; export * from './output'; diff --git a/packages/core/src/plugins/output/index.ts b/packages/core/src/plugins/output/index.ts index aece3f6..4628db7 100644 --- a/packages/core/src/plugins/output/index.ts +++ b/packages/core/src/plugins/output/index.ts @@ -4,6 +4,7 @@ export * from './types'; // Output handler implementations export { ConsoleOutputHandler } from './console-handler'; export { WebhookOutputHandler } from './webhook-handler'; +export { TempFileOutputHandler } from './temp-file-handler'; // Output manager export { outputManager, output, outputTo } from './output-manager'; @@ -36,6 +37,20 @@ export function registerWebhookOutput(config: import('./types').WebhookOutputCon return outputManager; } +/** + * Convenience function: Create and register a Temp File output handler + * @param config Temp file output handler configuration + * @returns Output manager instance + */ +export function registerTempFileOutput(config?: import('./types').TempFileOutputConfig) { + const { TempFileOutputHandler } = require('./temp-file-handler'); + const handler = new TempFileOutputHandler(config); + const { outputManager } = require('./output-manager'); + const name = 'temp-file_' + Date.now(); + outputManager.registerHandler(name, handler); + return outputManager; +} + /** * Convenience function: Register output handlers in batch * @param configs Output handler configuration array diff --git a/packages/core/src/plugins/output/output-manager.ts b/packages/core/src/plugins/output/output-manager.ts index e4c97a8..af08011 100644 --- a/packages/core/src/plugins/output/output-manager.ts +++ b/packages/core/src/plugins/output/output-manager.ts @@ -1,6 +1,7 @@ import { OutputHandler, OutputOptions, OutputHandlerConfig } from './types'; import { ConsoleOutputHandler } from './console-handler'; import { WebhookOutputHandler } from './webhook-handler'; +import { TempFileOutputHandler } from './temp-file-handler'; /** * Output manager @@ -51,6 +52,9 @@ class OutputManager { case 'webhook': return new WebhookOutputHandler(config.config as any); + case 'temp-file': + return new TempFileOutputHandler(config.config as any); + // Reserved for other output handler types // case 'websocket': // return new WebSocketOutputHandler(config.config as any); diff --git a/packages/core/src/plugins/output/temp-file-handler.ts b/packages/core/src/plugins/output/temp-file-handler.ts new file mode 100644 index 0000000..fa1fea5 --- /dev/null +++ b/packages/core/src/plugins/output/temp-file-handler.ts @@ -0,0 +1,140 @@ +import { OutputHandler, OutputOptions } from './types'; +import { writeFileSync, existsSync, mkdirSync } from 'fs'; +import { join } from 'path'; +import { tmpdir } from 'os'; + +/** + * Temp file output handler configuration + */ +export interface TempFileOutputConfig { + /** + * Subdirectory under system temp directory (default: 'claude-code-router') + */ + subdirectory?: string; + + /** + * File extension (default: 'json') + */ + extension?: string; + + /** + * Whether to include timestamp in filename (default: true) + */ + includeTimestamp?: boolean; + + /** + * Custom prefix for temp files (default: 'session') + */ + prefix?: string; +} + +/** + * Temp file output handler + * Writes data to temporary files in system temp directory + */ +export class TempFileOutputHandler implements OutputHandler { + type = 'temp-file' as const; + private config: TempFileOutputConfig; + private baseDir: string; + + constructor(config: TempFileOutputConfig = {}) { + this.config = { + subdirectory: 'claude-code-router', + extension: 'json', + includeTimestamp: true, + prefix: 'session', + ...config + }; + + // Use system temp directory + const systemTempDir = tmpdir(); + this.baseDir = join(systemTempDir, this.config.subdirectory!); + + // Ensure directory exists + this.ensureDir(); + } + + /** + * Ensure directory exists + */ + private ensureDir(): void { + try { + if (!existsSync(this.baseDir)) { + mkdirSync(this.baseDir, { recursive: true }); + } + } catch (error) { + // Silently fail + } + } + + /** + * Extract session ID from user_id string + * Format: "user_..._session_" + */ + private extractSessionId(userId: string): string | null { + try { + const match = userId.match(/_session_([a-f0-9-]+)/i); + return match ? match[1] : null; + } catch { + return null; + } + } + + /** + * Get file path for temp file + */ + private getFilePath(sessionId: string): string { + const prefix = this.config.prefix || 'session'; + const ext = this.config.extension ? `.${this.config.extension}` : ''; + + let filename: string; + if (this.config.includeTimestamp) { + // Include timestamp in filename: prefix-sessionId-timestamp.ext + const timestamp = Date.now(); + filename = `${prefix}-${sessionId}-${timestamp}${ext}`; + } else { + // Simple filename: prefix-sessionId.ext + filename = `${prefix}-${sessionId}${ext}`; + } + + return join(this.baseDir, filename); + } + + /** + * Output data to temp file + */ + async output(data: any, options: OutputOptions = {}): Promise { + try { + // Extract session ID from metadata + const sessionId = options.metadata?.sessionId; + + if (!sessionId) { + // No session ID, skip output + return false; + } + + // Prepare output data + const outputData = { + ...data, + timestamp: Date.now(), + sessionId + }; + + // Write to file + const filePath = this.getFilePath(sessionId); + writeFileSync(filePath, JSON.stringify(outputData, null, 2), 'utf-8'); + + return true; + } catch (error) { + // Silently fail to avoid disrupting main flow + return false; + } + } + + /** + * Get the base directory where temp files are stored + */ + getBaseDir(): string { + return this.baseDir; + } +} diff --git a/packages/core/src/plugins/output/types.ts b/packages/core/src/plugins/output/types.ts index c8d9a6f..a6e65f7 100644 --- a/packages/core/src/plugins/output/types.ts +++ b/packages/core/src/plugins/output/types.ts @@ -135,6 +135,31 @@ export interface WebSocketOutputConfig { }; } +/** + * Temp file output handler configuration + */ +export interface TempFileOutputConfig { + /** + * Subdirectory under system temp directory (default: 'claude-code-router') + */ + subdirectory?: string; + + /** + * File extension (default: 'json') + */ + extension?: string; + + /** + * Whether to include timestamp in filename (default: true) + */ + includeTimestamp?: boolean; + + /** + * Custom prefix for temp files (default: 'session') + */ + prefix?: string; +} + /** * Output handler registration configuration */ @@ -142,7 +167,7 @@ export interface OutputHandlerConfig { /** * Output handler type */ - type: 'console' | 'webhook' | 'websocket'; + type: 'console' | 'webhook' | 'websocket' | 'temp-file'; /** * Whether enabled @@ -152,5 +177,5 @@ export interface OutputHandlerConfig { /** * Configuration options */ - config?: ConsoleOutputConfig | WebhookOutputConfig | WebSocketOutputConfig; + config?: ConsoleOutputConfig | WebhookOutputConfig | WebSocketOutputConfig | TempFileOutputConfig; } diff --git a/packages/core/src/plugins/token-speed.ts b/packages/core/src/plugins/token-speed.ts index 91c0eb2..60ab745 100644 --- a/packages/core/src/plugins/token-speed.ts +++ b/packages/core/src/plugins/token-speed.ts @@ -9,26 +9,27 @@ import { ITokenizer, TokenizerConfig } from '../types/tokenizer'; */ interface TokenStats { requestId: string; + sessionId?: string; startTime: number; firstTokenTime?: number; lastTokenTime: number; tokenCount: number; tokensPerSecond: number; timeToFirstToken?: number; - contentBlocks: { - index: number; - tokenCount: number; - speed: number; - }[]; + stream: boolean; // Whether this is a streaming request + tokenTimestamps: number[]; // Store timestamps of each token for per-second calculation } /** * Plugin options */ interface TokenSpeedOptions extends CCRPluginOptions { - logInterval?: number; // Log every N tokens - enableCrossRequestStats?: boolean; // Enable cross-request statistics - statsWindow?: number; // Statistics window size (last N requests) + /** + * Reporter type(s) to use for output + * Can be a single type or an array of types: 'console' | 'temp-file' | 'webhook' + * Default: ['console', 'temp-file'] + */ + reporter?: string | string[]; /** * Output handler configurations @@ -48,18 +49,6 @@ const requestStats = new Map(); // Cache tokenizers by provider and model to avoid repeated initialization const tokenizerCache = new Map(); -// Cross-request statistics -const globalStats = { - totalRequests: 0, - totalTokens: 0, - totalTime: 0, - avgTokensPerSecond: 0, - minTokensPerSecond: Infinity, - maxTokensPerSecond: 0, - avgTimeToFirstToken: 0, - allSpeeds: [] as number[] // Used for calculating percentiles -}; - /** * Token speed measurement plugin */ @@ -71,25 +60,50 @@ export const tokenSpeedPlugin: CCRPlugin = { // Use fp() to break encapsulation and apply hooks globally register: fp(async (fastify, options: TokenSpeedOptions) => { const opts = { - logInterval: 10, - enableCrossRequestStats: true, - statsWindow: 100, + reporter: ['console', 'temp-file'], ...options }; - // Initialize output handlers + // Normalize reporter to array + const reporters = Array.isArray(opts.reporter) ? opts.reporter : [opts.reporter]; + + // Initialize output handlers based on reporters if not explicitly configured if (opts.outputHandlers && opts.outputHandlers.length > 0) { outputManager.registerHandlers(opts.outputHandlers); } else { - // Default to console output if no handlers configured - outputManager.registerHandlers([{ - type: 'console', - enabled: true, - config: { - colors: true, - level: 'log' + // Auto-register handlers based on reporter types + const handlersToRegister: OutputHandlerConfig[] = []; + + for (const reporter of reporters) { + if (reporter === 'console') { + handlersToRegister.push({ + type: 'console', + enabled: true, + config: { + colors: true, + level: 'log' + } + }); + } else if (reporter === 'temp-file') { + handlersToRegister.push({ + type: 'temp-file', + enabled: true, + config: { + subdirectory: 'claude-code-router', + extension: 'json', + includeTimestamp: true, + prefix: 'session' + } + }); + } else if (reporter === 'webhook') { + // Webhook requires explicit config, skip auto-registration + console.warn(`[TokenSpeedPlugin] Webhook reporter requires explicit configuration in outputHandlers`); } - }]); + } + + if (handlersToRegister.length > 0) { + outputManager.registerHandlers(handlersToRegister); + } } // Set default output options @@ -144,182 +158,243 @@ export const tokenSpeedPlugin: CCRPlugin = { } }; - // Add onSend hook to intercept streaming responses - fastify.addHook('onSend', async (request, reply, payload) => { - // Only handle streaming responses - if (!(payload instanceof ReadableStream)) { - return payload; - } + // Add onRequest hook to capture actual request start time (before processing) + fastify.addHook('onRequest', async (request) => { + (request as any).requestStartTime = performance.now(); + }); + // Add onSend hook to intercept both streaming and non-streaming responses + fastify.addHook('onSend', async (request, _reply, payload) => { const requestId = (request as any).id || Date.now().toString(); - const startTime = Date.now(); + const startTime = (request as any).requestStartTime || performance.now(); - // Initialize statistics - requestStats.set(requestId, { - requestId, - startTime, - lastTokenTime: startTime, - tokenCount: 0, - tokensPerSecond: 0, - contentBlocks: [] - }); + // Extract session ID from request body metadata + let sessionId: string | undefined; + try { + const userId = (request.body as any)?.metadata?.user_id; + if (userId && typeof userId === 'string') { + const match = userId.match(/_session_([a-f0-9-]+)/i); + sessionId = match ? match[1] : undefined; + } + } catch (error) { + // Ignore errors extracting session ID + } // Get tokenizer for this specific request const tokenizer = await getTokenizerForRequest(request); - // Tee the stream: one for stats, one for the client - const [originalStream, statsStream] = payload.tee(); + // Handle streaming responses + if (payload instanceof ReadableStream) { + // Mark this request as streaming + requestStats.set(requestId, { + requestId, + sessionId, + startTime, + lastTokenTime: startTime, + tokenCount: 0, + tokensPerSecond: 0, + tokenTimestamps: [], + stream: true + }); - // Process stats in background - const processStats = async () => { - let currentBlockIndex = -1; - let blockStartTime = 0; - let blockTokenCount = 0; + // Tee the stream: one for stats, one for the client + const [originalStream, statsStream] = payload.tee(); - try { - // Decode byte stream to text, then parse SSE events - const eventStream = statsStream - .pipeThrough(new TextDecoderStream()) - .pipeThrough(new SSEParserTransform()); - const reader = eventStream.getReader(); + // Process stats in background + const processStats = async () => { + let outputTimer: NodeJS.Timeout | null = null; - while (true) { - const { done, value } = await reader.read(); - if (done) break; - - const data = value; + // Output stats function - calculate current speed using sliding window + const doOutput = async (isFinal: boolean) => { const stats = requestStats.get(requestId); - if (!stats) continue; + if (!stats) return; - // Detect content_block_start event - if (data.event === 'content_block_start' && data.data?.content_block?.type === 'text') { - currentBlockIndex = data.data.index; - blockStartTime = Date.now(); - blockTokenCount = 0; - } + const now = performance.now(); - // Detect content_block_delta event (incremental tokens) - if (data.event === 'content_block_delta' && data.data?.delta?.type === 'text_delta') { - const text = data.data.delta.text; - const tokenCount = tokenizer - ? (tokenizer.encodeText ? tokenizer.encodeText(text).length : estimateTokens(text)) - : estimateTokens(text); - - stats.tokenCount += tokenCount; - stats.lastTokenTime = Date.now(); - - // Record first token time - if (!stats.firstTokenTime) { - stats.firstTokenTime = stats.lastTokenTime; - stats.timeToFirstToken = stats.firstTokenTime - stats.startTime; - } - - // Calculate current block token count - if (currentBlockIndex >= 0) { - blockTokenCount += tokenCount; - } - - // Calculate speed - const elapsed = (stats.lastTokenTime - stats.startTime) / 1000; - stats.tokensPerSecond = stats.tokenCount / elapsed; - - // Log periodically - if (stats.tokenCount % opts.logInterval === 0) { - await outputStats(stats, opts.outputOptions); + if (!isFinal) { + // For streaming output, use sliding window: count tokens in last 1 second + const oneSecondAgo = now - 1000; + stats.tokenTimestamps = stats.tokenTimestamps.filter(ts => ts > oneSecondAgo); + stats.tokensPerSecond = stats.tokenTimestamps.length; + } else { + // For final output, use average speed over entire request + const duration = (stats.lastTokenTime - stats.startTime) / 1000; // seconds + if (duration > 0) { + stats.tokensPerSecond = Math.round(stats.tokenCount / duration); } } - // Detect content_block_stop event - if (data.event === 'content_block_stop' && currentBlockIndex >= 0) { - const blockElapsed = (Date.now() - blockStartTime) / 1000; - const blockSpeed = blockElapsed > 0 ? blockTokenCount / blockElapsed : 0; + await outputStats(stats, reporters, opts.outputOptions, isFinal).catch(err => { + fastify.log?.warn(`Failed to output streaming stats: ${err.message}`); + }); + }; - stats.contentBlocks.push({ - index: currentBlockIndex, - tokenCount: blockTokenCount, - speed: blockSpeed - }); + try { + // Decode byte stream to text, then parse SSE events + const eventStream = statsStream + .pipeThrough(new TextDecoderStream()) + .pipeThrough(new SSEParserTransform()); + const reader = eventStream.getReader(); - currentBlockIndex = -1; + // Start timer immediately - output every 1 second + outputTimer = setInterval(async () => { + const stats = requestStats.get(requestId); + if (stats) { + await doOutput(false); + } + }, 1000); + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + const data = value; + const stats = requestStats.get(requestId); + if (!stats) continue; + + const now = performance.now(); + + // Record first token time when we receive any content-related event + // This includes: content_block_start, content_block_delta, text_block + if (!stats.firstTokenTime && ( + data.event === 'content_block_start' || + data.event === 'content_block_delta' || + data.event === 'text_block' || + data.event === 'content_block' + )) { + stats.firstTokenTime = now; + stats.timeToFirstToken = Math.round(now - stats.startTime); + } + + // Detect content_block_delta event (incremental tokens) + // Support multiple delta types: text_delta, input_json_delta, thinking_delta + if (data.event === 'content_block_delta' && data.data?.delta) { + const deltaType = data.data.delta.type; + let text = ''; + + // Extract text based on delta type + if (deltaType === 'text_delta') { + text = data.data.delta.text || ''; + } else if (deltaType === 'input_json_delta') { + text = data.data.delta.partial_json || ''; + } else if (deltaType === 'thinking_delta') { + text = data.data.delta.thinking || ''; + } + + // Calculate tokens if we have text content + if (text) { + const tokenCount = tokenizer + ? (tokenizer.encodeText ? tokenizer.encodeText(text).length : estimateTokens(text)) + : estimateTokens(text); + + stats.tokenCount += tokenCount; + stats.lastTokenTime = now; + + // Record timestamps for each token (for sliding window calculation) + for (let i = 0; i < tokenCount; i++) { + stats.tokenTimestamps.push(now); + } + } + } + + // Output final statistics when message ends + if (data.event === 'message_stop') { + // Clear timer + if (outputTimer) { + clearInterval(outputTimer); + outputTimer = null; + } + + await doOutput(true); + + requestStats.delete(requestId); + } } - - // Output final statistics when message ends - if (data.event === 'message_stop') { - // Update global statistics - if (opts.enableCrossRequestStats) { - updateGlobalStats(stats, opts.statsWindow); - } - - await outputStats(stats, opts.outputOptions, true); - - if (opts.enableCrossRequestStats) { - await outputGlobalStats(opts.outputOptions); - } - - requestStats.delete(requestId); + } catch (error: any) { + // Clean up timer on error + if (outputTimer) { + clearInterval(outputTimer); + } + if (error.name !== 'AbortError' && error.code !== 'ERR_STREAM_PREMATURE_CLOSE') { + fastify.log?.warn(`Error processing token stats: ${error.message}`); } } - } catch (error: any) { - console.error(error); - if (error.name !== 'AbortError' && error.code !== 'ERR_STREAM_PREMATURE_CLOSE') { - fastify.log?.warn(`Error processing token stats: ${error.message}`); + }; + + // Start background processing without blocking + processStats().catch((error) => { + console.log(error); + fastify.log?.warn(`Background stats processing failed: ${error.message}`); + }); + + // Return original stream to client + return originalStream; + } + + // Handle non-streaming responses + // Try to extract token count from the response payload + const endTime = performance.now(); + let tokenCount = 0; + + // Payload should be a string or object for non-streaming responses + if (payload && typeof payload === 'string') { + try { + const response = JSON.parse(payload); + + // Prefer usage.output_tokens if available (most accurate) + if (response.usage?.output_tokens) { + tokenCount = response.usage.output_tokens; + } else { + // Fallback: calculate from content + const content = response.content || response.message?.content || ''; + + if (tokenizer) { + if (Array.isArray(content)) { + tokenCount = content.reduce((sum: number, block: any) => { + if (block.type === 'text') { + const text = block.text || ''; + return sum + (tokenizer.encodeText ? tokenizer.encodeText(text).length : estimateTokens(text)); + } + return sum; + }, 0); + } else if (typeof content === 'string') { + tokenCount = tokenizer.encodeText ? tokenizer.encodeText(content).length : estimateTokens(content); + } + } else { + const text = Array.isArray(content) ? content.map((c: any) => c.text).join('') : content; + tokenCount = estimateTokens(text); + } } + } catch (error) { + // Could not parse or extract tokens } - }; + } - // Start background processing without blocking - processStats().catch((error) => { - console.log(error); - fastify.log?.warn(`Background stats processing failed: ${error.message}`); - }); + // Only output stats if we found tokens + if (tokenCount > 0) { + const duration = (endTime - startTime) / 1000; // seconds - // Return original stream to client - return originalStream; + const stats: TokenStats = { + requestId, + sessionId, + startTime, + lastTokenTime: endTime, + tokenCount, + tokensPerSecond: duration > 0 ? Math.round(tokenCount / duration) : 0, + timeToFirstToken: Math.round(endTime - startTime), + stream: false, + tokenTimestamps: [] + }; + + await outputStats(stats, reporters, opts.outputOptions, true); + } + + // Return payload as-is + return payload; }); }), }; -/** - * Update global statistics - */ -function updateGlobalStats(stats: TokenStats, windowSize: number) { - globalStats.totalRequests++; - globalStats.totalTokens += stats.tokenCount; - globalStats.totalTime += (stats.lastTokenTime - stats.startTime) / 1000; - - if (stats.tokensPerSecond < globalStats.minTokensPerSecond) { - globalStats.minTokensPerSecond = stats.tokensPerSecond; - } - if (stats.tokensPerSecond > globalStats.maxTokensPerSecond) { - globalStats.maxTokensPerSecond = stats.tokensPerSecond; - } - - if (stats.timeToFirstToken) { - globalStats.avgTimeToFirstToken = - (globalStats.avgTimeToFirstToken * (globalStats.totalRequests - 1) + stats.timeToFirstToken) / - globalStats.totalRequests; - } - - globalStats.allSpeeds.push(stats.tokensPerSecond); - - // Maintain window size - if (globalStats.allSpeeds.length > windowSize) { - globalStats.allSpeeds.shift(); - } - - globalStats.avgTokensPerSecond = globalStats.totalTokens / globalStats.totalTime; -} - -/** - * Calculate percentile - */ -function calculatePercentile(data: number[], percentile: number): number { - if (data.length === 0) return 0; - const sorted = [...data].sort((a, b) => a - b); - const index = Math.ceil((percentile / 100) * sorted.length) - 1; - return sorted[index]; -} - /** * Estimate token count (fallback method) */ @@ -333,63 +408,39 @@ function estimateTokens(text: string): number { /** * Output single request statistics */ -async function outputStats(stats: TokenStats, options?: OutputOptions, isFinal = false) { +async function outputStats( + stats: TokenStats, + reporters: string[], + options?: OutputOptions, + isFinal = false +) { const prefix = isFinal ? '[Token Speed Final]' : '[Token Speed]'; - // Calculate average speed of each block - const avgBlockSpeed = stats.contentBlocks.length > 0 - ? stats.contentBlocks.reduce((sum, b) => sum + b.speed, 0) / stats.contentBlocks.length - : 0; - const logData = { requestId: stats.requestId.substring(0, 8), + sessionId: stats.sessionId, + stream: stats.stream, tokenCount: stats.tokenCount, - tokensPerSecond: stats.tokensPerSecond.toFixed(2), + tokensPerSecond: stats.tokensPerSecond, timeToFirstToken: stats.timeToFirstToken ? `${stats.timeToFirstToken}ms` : 'N/A', duration: `${((stats.lastTokenTime - stats.startTime) / 1000).toFixed(2)}s`, - contentBlocks: stats.contentBlocks.length, - avgBlockSpeed: avgBlockSpeed.toFixed(2), - ...(isFinal && stats.contentBlocks.length > 1 ? { - blocks: stats.contentBlocks.map(b => ({ - index: b.index, - tokenCount: b.tokenCount, - speed: b.speed.toFixed(2) - })) - } : {}) + timestamp: Date.now() }; - // Output through output manager - await outputManager.output(logData, { + const outputOptions = { prefix, + metadata: { + sessionId: stats.sessionId + }, ...options - }); -} - -/** - * Output global statistics - */ -async function outputGlobalStats(options?: OutputOptions) { - const p50 = calculatePercentile(globalStats.allSpeeds, 50); - const p95 = calculatePercentile(globalStats.allSpeeds, 95); - const p99 = calculatePercentile(globalStats.allSpeeds, 99); - - const logData = { - totalRequests: globalStats.totalRequests, - totalTokens: globalStats.totalTokens, - avgTokensPerSecond: globalStats.avgTokensPerSecond.toFixed(2), - minSpeed: globalStats.minTokensPerSecond === Infinity ? 0 : globalStats.minTokensPerSecond.toFixed(2), - maxSpeed: globalStats.maxTokensPerSecond.toFixed(2), - avgTimeToFirstToken: `${globalStats.avgTimeToFirstToken.toFixed(0)}ms`, - percentiles: { - p50: p50.toFixed(2), - p95: p95.toFixed(2), - p99: p99.toFixed(2) - } }; - // Output through output manager - await outputManager.output(logData, { - prefix: '[Token Speed Global Stats]', - ...options - }); + // Output to each specified reporter type + for (const reporter of reporters) { + try { + await outputManager.outputToType(reporter, logData, outputOptions); + } catch (error) { + console.error(`[TokenSpeedPlugin] Failed to output to ${reporter}:`, error); + } + } } diff --git a/packages/core/src/server.ts b/packages/core/src/server.ts index c02422e..9dc34bc 100644 --- a/packages/core/src/server.ts +++ b/packages/core/src/server.ts @@ -270,5 +270,5 @@ export { ConfigService } from "./services/config"; export { ProviderService } from "./services/provider"; export { TransformerService } from "./services/transformer"; export { TokenizerService } from "./services/tokenizer"; -export { pluginManager, tokenSpeedPlugin, CCRPlugin, CCRPluginOptions, PluginMetadata } from "./plugins"; +export { pluginManager, tokenSpeedPlugin, getTokenSpeedStats, getGlobalTokenSpeedStats, CCRPlugin, CCRPluginOptions, PluginMetadata } from "./plugins"; export { SSEParserTransform, SSESerializerTransform, rewriteStream } from "./utils/sse"; diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts index 126cf47..9f8b8a7 100644 --- a/packages/server/src/index.ts +++ b/packages/server/src/index.ts @@ -69,6 +69,12 @@ async function registerPluginsFromConfig(serverInstance: any, config: any): Prom case 'token-speed': pluginManager.registerPlugin(tokenSpeedPlugin, { enabled, + outputHandlers: [ + { + type: 'temp-file', + enabled: true + } + ], ...options }); break; diff --git a/packages/server/src/middleware/auth.ts b/packages/server/src/middleware/auth.ts index 2ee3b65..0b0ba22 100644 --- a/packages/server/src/middleware/auth.ts +++ b/packages/server/src/middleware/auth.ts @@ -4,7 +4,8 @@ export const apiKeyAuth = (config: any) => async (req: FastifyRequest, reply: FastifyReply, done: () => void) => { // Public endpoints that don't require authentication - if (["/", "/health"].includes(req.url) || req.url.startsWith("/ui")) { + const publicPaths = ["/", "/health"]; + if (publicPaths.includes(req.url) || req.url.startsWith("/ui")) { return done(); } diff --git a/packages/server/src/types.d.ts b/packages/server/src/types.d.ts index de4e0f8..8e2f7d0 100644 --- a/packages/server/src/types.d.ts +++ b/packages/server/src/types.d.ts @@ -118,4 +118,46 @@ declare module "@musistudio/llms" { clearCache(): void; dispose(): void; } + + // Token speed statistics types + export interface TokenStats { + requestId: string; + startTime: number; + firstTokenTime?: number; + lastTokenTime: number; + tokenCount: number; + tokensPerSecond: number; + timeToFirstToken?: number; + contentBlocks: { + index: number; + tokenCount: number; + speed: number; + }[]; + } + + export function getTokenSpeedStats(): { + current: TokenStats | null; + global: { + totalRequests: number; + totalTokens: number; + totalTime: number; + avgTokensPerSecond: number; + minTokensPerSecond: number; + maxTokensPerSecond: number; + avgTimeToFirstToken: number; + allSpeeds: number[]; + }; + lastUpdate: number; + }; + + export function getGlobalTokenSpeedStats(): { + totalRequests: number; + totalTokens: number; + totalTime: number; + avgTokensPerSecond: number; + minTokensPerSecond: number; + maxTokensPerSecond: number; + avgTimeToFirstToken: number; + allSpeeds: number[]; + }; } diff --git a/packages/ui/src/components/StatusLineConfigDialog.tsx b/packages/ui/src/components/StatusLineConfigDialog.tsx index 3dd061f..a7678bb 100644 --- a/packages/ui/src/components/StatusLineConfigDialog.tsx +++ b/packages/ui/src/components/StatusLineConfigDialog.tsx @@ -51,6 +51,7 @@ const MODULE_TYPES = [ { label: "gitBranch", value: "gitBranch" }, { label: "model", value: "model" }, { label: "usage", value: "usage" }, + { label: "speed", value: "speed" }, { label: "script", value: "script" }, ]; @@ -936,6 +937,14 @@ export function StatusLineConfigDialog({ color: "bright_magenta", }; break; + case "speed": + newModule = { + type: "speed", + icon: "⚡", + text: "{{tokenSpeed}}", + color: "bright_green", + }; + break; case "script": newModule = { type: "script", diff --git a/scripts/build-cli.js b/scripts/build-cli.js index 2392fe2..04b4189 100644 --- a/scripts/build-cli.js +++ b/scripts/build-cli.js @@ -46,7 +46,7 @@ try { // Step 4: Build the CLI application console.log('Building CLI application...'); - execSync('esbuild src/cli.ts --bundle --platform=node --outfile=dist/cli.js', { + execSync('esbuild src/cli.ts --bundle --platform=node --minify --tree-shaking=true --outfile=dist/cli.js', { stdio: 'inherit', cwd: cliDir });