From bac32fe8f6f4524d394b2be6a035133f5599e302 Mon Sep 17 00:00:00 2001 From: kurihada Date: Thu, 18 Dec 2025 16:11:00 +0800 Subject: [PATCH] =?UTF-8?q?feat(core):=20=E5=AE=9E=E7=8E=B0=20Token=20?= =?UTF-8?q?=E6=B6=88=E8=80=97=E7=BB=9F=E8=AE=A1=E7=B3=BB=E7=BB=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 扩展 SessionStats schema 添加 token 统计字段 - 添加 TokenUsageInfo 类型和 ChatResult.usage 字段 - AgentMessageHandler 从 AI SDK response 提取 usage - AgentExecutor 返回 usage 到执行结果 - 新增 TokenStatsManager 管理统计: - updateSessionStats: 更新会话 token 统计 - mergeChildSessionStats: 合并子会话统计到父会话 - getSessionStats/getProjectStats: 查询统计 - Agent.chat() 完成后自动更新统计 - Task 工具完成后合并子会话统计 - 新增 REST API: /api/stats/sessions/:id, /api/stats/projects/:id - 添加 TokenStatsManager 单元测试 (12 tests) --- packages/core/src/agent/executor.ts | 33 +- packages/core/src/agent/types.ts | 4 +- .../core/src/core/agent-message-handler.ts | 49 ++- packages/core/src/core/agent.ts | 16 +- packages/core/src/index.ts | 4 +- packages/core/src/session/index.ts | 7 + packages/core/src/session/project-manager.ts | 12 + packages/core/src/session/session-store.ts | 2 + packages/core/src/session/storage/index.ts | 4 +- packages/core/src/session/storage/session.ts | 35 +- .../core/src/session/token-stats-manager.ts | 287 ++++++++++++++ packages/core/src/tools/task/task.ts | 18 +- packages/core/src/types/index.ts | 18 + .../unit/session/token-stats-manager.test.ts | 372 ++++++++++++++++++ packages/server/src/index.ts | 3 +- packages/server/src/routes/index.ts | 1 + packages/server/src/routes/stats.ts | 216 ++++++++++ 17 files changed, 1054 insertions(+), 27 deletions(-) create mode 100644 packages/core/src/session/token-stats-manager.ts create mode 100644 packages/core/tests/unit/session/token-stats-manager.test.ts create mode 100644 packages/server/src/routes/stats.ts diff --git a/packages/core/src/agent/executor.ts b/packages/core/src/agent/executor.ts index 42cfbe0..84e614e 100644 --- a/packages/core/src/agent/executor.ts +++ b/packages/core/src/agent/executor.ts @@ -6,7 +6,7 @@ import { type Tool as AITool, type LanguageModel, } from 'ai'; -import type { Tool, ToolResult, AgentConfig, ContentBlock } from '../types/index.js'; +import type { Tool, ToolResult, AgentConfig, ContentBlock, TokenUsageInfo } from '../types/index.js'; import { buildZodSchema } from '../types/index.js'; import { ToolRegistry } from '../tools/registry.js'; import type { @@ -95,6 +95,7 @@ export class AgentExecutor { let fullResponse = ''; let steps = 0; + let usage: TokenUsageInfo | undefined; // 工具调用时间追踪（用于计算持续时间） const toolStartTimes = new Map(); @@ -200,7 +201,9 @@ export class AgentExecutor { } } - await result.response; + const response = await result.response; + // 提取 usage 信息 + usage = this.extractUsage(response); } else { // 非流式模式 const result = await generateText({ @@ -214,6 +217,8 @@ export class AgentExecutor { fullResponse = result.text; steps = result.steps.length; + // 提取 usage 信息 + usage = this.extractUsage(result.response); } return { @@ -221,6 +226,7 @@ export class AgentExecutor { text: fullResponse, steps, sessionId: context.parentSessionId ?? 'standalone', + usage, }; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); @@ -416,4 +422,27 @@ export class AgentExecutor { return blocks; } + + /** + * 从 AI SDK 响应中提取 usage 信息 + */ + private extractUsage(response: unknown): TokenUsageInfo | undefined { + // AI SDK 的 response 对象包含 usage 字段 + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const resp = response as any; + const usage = resp?.usage; + + if (!usage) { + return undefined; + } + + return { + promptTokens: usage.promptTokens ?? 0, + completionTokens: usage.completionTokens ?? 0, + totalTokens: usage.totalTokens ?? (usage.promptTokens ?? 0) + (usage.completionTokens ?? 0), + // Anthropic API 特有的缓存字段 + cacheReadInputTokens: usage.cacheReadInputTokens, + cacheCreationInputTokens: usage.cacheCreationInputTokens, + }; + } } diff --git a/packages/core/src/agent/types.ts b/packages/core/src/agent/types.ts index 8db3223..68c7196 100644 --- a/packages/core/src/agent/types.ts +++ b/packages/core/src/agent/types.ts @@ -1,4 +1,4 @@ -import type { ProviderType } from '../types/index.js'; +import type { ProviderType, TokenUsageInfo } from '../types/index.js'; import type { PermissionAction, PermissionRule } from '../permission/types.js'; // 重新导出权限类型，方便外部使用 @@ -196,4 +196,6 @@ export interface AgentExecutionResult { sessionId: string; /** 错误信息 */ error?: string; + /** Token 使用统计 */ + usage?: TokenUsageInfo; } diff --git a/packages/core/src/core/agent-message-handler.ts b/packages/core/src/core/agent-message-handler.ts index ecd5179..a66283e 100644 --- a/packages/core/src/core/agent-message-handler.ts +++ b/packages/core/src/core/agent-message-handler.ts @@ -11,7 +11,7 @@ import { type Tool as AITool, type LanguageModel, } from 'ai'; -import type { ToolResult, UserInput, ContentBlock, ChatResult } from '../types/index.js'; +import type { ToolResult, UserInput, ContentBlock, ChatResult, TokenUsageInfo } from '../types/index.js'; import { CompressionManager, CompressionStatus, @@ -188,6 +188,18 @@ export class AgentMessageHandler { const response = await result.response; responseMessages = response.messages as ModelMessage[]; + + // 提取 usage 信息 + const usage = this.extractUsage(response); + + // 将完整的响应消息添加到历史 + this.conversationHistory.push(...responseMessages); + + return { + text: fullResponse, + messages: responseMessages, + usage, + }; } catch (error) { if (error instanceof Error && (error.name === 'AbortError' || abortSignal?.aborted)) { onStream?.('\n[已取消]\n'); @@ -201,14 +213,6 @@ export class AgentMessageHandler { } throw error; } - - // 将完整的响应消息添加到历史 - this.conversationHistory.push(...responseMessages); - - return { - text: fullResponse, - messages: responseMessages, - }; } /** @@ -320,12 +324,39 @@ export class AgentMessageHandler { const fullResponse = result.text; const responseMessages = result.response.messages as ModelMessage[]; + // 提取 usage 信息 + const usage = this.extractUsage(result.response); + // 将完整的响应消息添加到历史 this.conversationHistory.push(...responseMessages); return { text: fullResponse, messages: responseMessages, + usage, + }; + } + + /** + * 从 AI SDK 响应中提取 usage 信息 + */ + private extractUsage(response: unknown): TokenUsageInfo | undefined { + // AI SDK 的 response 对象包含 usage 字段 + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const resp = response as any; + const usage = resp?.usage; + + if (!usage) { + return undefined; + } + + return { + promptTokens: usage.promptTokens ?? 0, + completionTokens: usage.completionTokens ?? 0, + totalTokens: usage.totalTokens ?? (usage.promptTokens ?? 0) + (usage.completionTokens ?? 0), + // Anthropic API 特有的缓存字段 + cacheReadInputTokens: usage.cacheReadInputTokens, + cacheCreationInputTokens: usage.cacheCreationInputTokens, }; } diff --git a/packages/core/src/core/agent.ts b/packages/core/src/core/agent.ts index 838139a..88f785d 100644 --- a/packages/core/src/core/agent.ts +++ b/packages/core/src/core/agent.ts @@ -4,9 +4,9 @@ */ import type { LanguageModel } from 'ai'; -import type { Tool, AgentConfig, UserInput, Message, ChatResult } from '../types/index.js'; +import type { Tool, AgentConfig, UserInput, Message, ChatResult, TokenUsageInfo } from '../types/index.js'; import { ToolRegistry } from '../tools/registry.js'; -import { SessionManager } from '../session/index.js'; +import { SessionManager, tokenStatsManager } from '../session/index.js'; import { CompressionManager, type TokenUsage, @@ -256,6 +256,18 @@ export class Agent { ); } + // 更新 Token 统计 + if (result.usage && this.sessionManager) { + const session = this.sessionManager.getSession(); + if (session) { + await tokenStatsManager.updateSessionStats( + session.projectId, + session.id, + result.usage + ); + } + } + // 自动压缩 await this.messageHandler.autoCompress(onStream); diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 77831bf..e645638 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -29,8 +29,8 @@ export type { DetailedCompressionResult, } from './context/index.js'; // Session - 新的三层存储结构 -export { SessionManager } from './session/index.js'; -export type { SessionData, SessionSummary, ProjectMetadata } from './session/index.js'; +export { SessionManager, tokenStatsManager, TokenStatsManager } from './session/index.js'; +export type { SessionData, SessionSummary, ProjectMetadata, SessionTokenStats, ProjectStats } from './session/index.js'; // Session Storage API export { diff --git a/packages/core/src/session/index.ts b/packages/core/src/session/index.ts index 7d764d8..c068ca1 100644 --- a/packages/core/src/session/index.ts +++ b/packages/core/src/session/index.ts @@ -77,3 +77,10 @@ export type { SessionData, SessionSummary, ProjectMetadata } from './manager.js' // 项目工具 export { getProjectId, isGitRepository } from './project.js'; + +// Token 统计管理器 +export { TokenStatsManager, tokenStatsManager } from './token-stats-manager.js'; +export type { SessionTokenStats } from './token-stats-manager.js'; + +// 项目管理器类型 +export type { ProjectStats } from './project-manager.js'; diff --git a/packages/core/src/session/project-manager.ts b/packages/core/src/session/project-manager.ts index cf787d3..0bba4be 100644 --- a/packages/core/src/session/project-manager.ts +++ b/packages/core/src/session/project-manager.ts @@ -6,6 +6,17 @@ import * as storage from './storage/index.js'; import { getProjectId, isGitRepository } from './project.js'; +/** + * 项目 Token 统计 + */ +export interface ProjectStats { + totalInputTokens: number; + totalOutputTokens: number; + totalTokens: number; + sessionCount: number; + updatedAt: number; +} + /** * 项目元数据 */ @@ -14,6 +25,7 @@ export interface ProjectMetadata { workdir: string; createdAt: string; isGitRepo: boolean; + stats?: ProjectStats; } /** diff --git a/packages/core/src/session/session-store.ts b/packages/core/src/session/session-store.ts index 3429052..37f6894 100644 --- a/packages/core/src/session/session-store.ts +++ b/packages/core/src/session/session-store.ts @@ -113,6 +113,8 @@ export class SessionStore { messageCount: session.messages.length, inputTokens: 0, outputTokens: 0, + totalTokens: 0, + updatedAt: Date.now(), }, }; diff --git a/packages/core/src/session/storage/index.ts b/packages/core/src/session/storage/index.ts index fedb4a5..ff0fa32 100644 --- a/packages/core/src/session/storage/index.ts +++ b/packages/core/src/session/storage/index.ts @@ -17,8 +17,8 @@ export { // Session storage export * as SessionStorage from './session.js'; -export type { SessionInfo } from './session.js'; -export { SessionInfoSchema } from './session.js'; +export type { SessionInfo, SessionStats } from './session.js'; +export { SessionInfoSchema, SessionStatsSchema, ChildSessionsTokensSchema } from './session.js'; // Message storage export * as MessageStorage from './message.js'; diff --git a/packages/core/src/session/storage/session.ts b/packages/core/src/session/storage/session.ts index faa9181..437fd2b 100644 --- a/packages/core/src/session/storage/session.ts +++ b/packages/core/src/session/storage/session.ts @@ -2,6 +2,33 @@ import { z } from 'zod'; import * as base from './base.js'; import { generateSessionId } from '../id.js'; +/** + * 子会话 Token 统计 Schema + */ +export const ChildSessionsTokensSchema = z.object({ + inputTokens: z.number(), + outputTokens: z.number(), + totalTokens: z.number(), +}); + +/** + * Session 统计信息 Schema + */ +export const SessionStatsSchema = z.object({ + messageCount: z.number(), + inputTokens: z.number(), + outputTokens: z.number(), + totalTokens: z.number(), + // 缓存相关（Anthropic API 支持） + cacheReadTokens: z.number().optional(), + cacheWriteTokens: z.number().optional(), + // 子会话统计（合并所有子会话的 token 消耗） + childSessionsTokens: ChildSessionsTokensSchema.optional(), + // 统计最后更新时间 + updatedAt: z.number(), +}); +export type SessionStats = z.infer; + /** * Session Info Schema */ @@ -16,13 +43,7 @@ export const SessionInfoSchema = z.object({ title: z.string().optional(), discoveredTools: z.array(z.string()).default([]), // 统计信息 - stats: z - .object({ - messageCount: z.number(), - inputTokens: z.number(), - outputTokens: z.number(), - }) - .optional(), + stats: SessionStatsSchema.optional(), }); export type SessionInfo = z.infer; diff --git a/packages/core/src/session/token-stats-manager.ts b/packages/core/src/session/token-stats-manager.ts new file mode 100644 index 0000000..02e52d9 --- /dev/null +++ b/packages/core/src/session/token-stats-manager.ts @@ -0,0 +1,287 @@ +/** + * Token 统计管理器 + * 负责 Session 和 Project 级别的 token 消耗统计 + */ + +import * as SessionStorage from './storage/session.js'; +import * as storage from './storage/index.js'; +import type { SessionStats } from './storage/session.js'; +import type { ProjectMetadata, ProjectStats } from './project-manager.js'; +import type { TokenUsageInfo } from '../types/index.js'; + +/** + * Session Token 统计详情 + */ +export interface SessionTokenStats { + /** 本会话自身的统计 */ + self: { + inputTokens: number; + outputTokens: number; + totalTokens: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; + }; + /** 子会话统计汇总 */ + children: { + inputTokens: number; + outputTokens: number; + totalTokens: number; + }; + /** 总计（自身 + 子会话） */ + total: { + inputTokens: number; + outputTokens: number; + totalTokens: number; + }; + /** 消息数量 */ + messageCount: number; + /** 最后更新时间 */ + updatedAt: number; +} + +/** + * Token 统计管理器 + */ +export class TokenStatsManager { + /** + * 更新 Session 级别统计 + * 将新的 token 使用累加到现有统计中 + */ + async updateSessionStats( + projectId: string, + sessionId: string, + usage: TokenUsageInfo + ): Promise { + const session = await SessionStorage.get(projectId, sessionId); + if (!session) { + return; + } + + const currentStats: SessionStats = session.stats || { + messageCount: 0, + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + updatedAt: Date.now(), + }; + + // 累加新的 token 使用 + const newStats: SessionStats = { + messageCount: currentStats.messageCount + 1, + inputTokens: currentStats.inputTokens + usage.promptTokens, + outputTokens: currentStats.outputTokens + usage.completionTokens, + totalTokens: currentStats.totalTokens + usage.totalTokens, + cacheReadTokens: (currentStats.cacheReadTokens || 0) + (usage.cacheReadInputTokens || 0), + cacheWriteTokens: (currentStats.cacheWriteTokens || 0) + (usage.cacheCreationInputTokens || 0), + childSessionsTokens: currentStats.childSessionsTokens, + updatedAt: Date.now(), + }; + + await SessionStorage.updateStats(projectId, sessionId, newStats); + } + + /** + * 合并子会话统计到父会话 + * 在子会话（Task 工具）完成后调用 + */ + async mergeChildSessionStats( + projectId: string, + parentSessionId: string, + childSessionId: string + ): Promise { + const parent = await SessionStorage.get(projectId, parentSessionId); + const child = await SessionStorage.get(projectId, childSessionId); + + if (!parent || !child || !child.stats) { + return; + } + + const parentStats: SessionStats = parent.stats || { + messageCount: 0, + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + updatedAt: Date.now(), + }; + + // 获取当前子会话累计 + const childTokens = parentStats.childSessionsTokens || { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + }; + + // 累加子会话的 token（包括子会话自身和其嵌套的子会话） + const childTotal = child.stats.inputTokens + (child.stats.childSessionsTokens?.inputTokens || 0); + const childOutputTotal = child.stats.outputTokens + (child.stats.childSessionsTokens?.outputTokens || 0); + const childTotalTokens = child.stats.totalTokens + (child.stats.childSessionsTokens?.totalTokens || 0); + + parentStats.childSessionsTokens = { + inputTokens: childTokens.inputTokens + childTotal, + outputTokens: childTokens.outputTokens + childOutputTotal, + totalTokens: childTokens.totalTokens + childTotalTokens, + }; + parentStats.updatedAt = Date.now(); + + await SessionStorage.updateStats(projectId, parentSessionId, parentStats); + } + + /** + * 获取 Session 完整统计 + * 包含自身统计、子会话统计和总计 + */ + async getSessionStats( + projectId: string, + sessionId: string + ): Promise { + const session = await SessionStorage.get(projectId, sessionId); + if (!session || !session.stats) { + return null; + } + + const stats = session.stats; + const childStats = stats.childSessionsTokens || { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + }; + + return { + self: { + inputTokens: stats.inputTokens, + outputTokens: stats.outputTokens, + totalTokens: stats.totalTokens, + cacheReadTokens: stats.cacheReadTokens, + cacheWriteTokens: stats.cacheWriteTokens, + }, + children: { + inputTokens: childStats.inputTokens, + outputTokens: childStats.outputTokens, + totalTokens: childStats.totalTokens, + }, + total: { + inputTokens: stats.inputTokens + childStats.inputTokens, + outputTokens: stats.outputTokens + childStats.outputTokens, + totalTokens: stats.totalTokens + childStats.totalTokens, + }, + messageCount: stats.messageCount, + updatedAt: stats.updatedAt, + }; + } + + /** + * 更新 Project 级别统计 + * 遍历项目下所有会话并汇总 token 消耗 + */ + async updateProjectStats(projectId: string): Promise { + const sessions = await SessionStorage.listByProject(projectId); + + let totalInput = 0; + let totalOutput = 0; + let totalTokens = 0; + + for (const session of sessions) { + // 跳过子会话（避免重复计算） + if (session.parentId) { + continue; + } + + if (session.stats) { + totalInput += session.stats.inputTokens; + totalOutput += session.stats.outputTokens; + totalTokens += session.stats.totalTokens; + + // 包含子会话统计 + if (session.stats.childSessionsTokens) { + totalInput += session.stats.childSessionsTokens.inputTokens; + totalOutput += session.stats.childSessionsTokens.outputTokens; + totalTokens += session.stats.childSessionsTokens.totalTokens; + } + } + } + + const projectStats: ProjectStats = { + totalInputTokens: totalInput, + totalOutputTokens: totalOutput, + totalTokens, + sessionCount: sessions.filter(s => !s.parentId).length, + updatedAt: Date.now(), + }; + + // 更新项目元数据 + try { + await storage.update(['project', projectId], (project: ProjectMetadata) => { + project.stats = projectStats; + }); + } catch { + // 项目可能不存在，忽略错误 + } + + return projectStats; + } + + /** + * 获取 Project 统计 + * 如果缓存的统计过期，会重新计算 + */ + async getProjectStats( + projectId: string, + forceRefresh = false + ): Promise { + try { + const project = await storage.read(['project', projectId]); + + // 如果需要强制刷新，或者没有缓存的统计 + if (forceRefresh || !project.stats) { + return await this.updateProjectStats(projectId); + } + + return project.stats; + } catch (e) { + if (e instanceof storage.StorageNotFoundError) { + return null; + } + throw e; + } + } + + /** + * 格式化 token 数量显示 + */ + formatTokens(tokens: number): string { + if (tokens >= 1000000) { + return `${(tokens / 1000000).toFixed(2)}M`; + } + if (tokens >= 1000) { + return `${(tokens / 1000).toFixed(1)}k`; + } + return `${tokens}`; + } + + /** + * 格式化 Session 统计为可读字符串 + */ + formatSessionStats(stats: SessionTokenStats): string { + const lines = [ + `Token 使用统计:`, + ` 输入: ${this.formatTokens(stats.total.inputTokens)}`, + ` 输出: ${this.formatTokens(stats.total.outputTokens)}`, + ` 总计: ${this.formatTokens(stats.total.totalTokens)}`, + ]; + + if (stats.children.totalTokens > 0) { + lines.push(` (其中子任务: ${this.formatTokens(stats.children.totalTokens)})`); + } + + if (stats.self.cacheReadTokens) { + lines.push(` 缓存读取: ${this.formatTokens(stats.self.cacheReadTokens)}`); + } + + lines.push(` 消息数: ${stats.messageCount}`); + + return lines.join('\n'); + } +} + +// 导出单例 +export const tokenStatsManager = new TokenStatsManager(); diff --git a/packages/core/src/tools/task/task.ts b/packages/core/src/tools/task/task.ts index 0195303..af9221f 100644 --- a/packages/core/src/tools/task/task.ts +++ b/packages/core/src/tools/task/task.ts @@ -4,7 +4,7 @@ import type { AgentConfig, ToolResult } from '../../types/index.js'; import type { ImageData } from '../../agent/types.js'; import { agentRegistry, AgentExecutor, agentEventEmitter } from '../../agent/index.js'; import { toolRegistry } from '../registry.js'; -import { SessionManager } from '../../session/index.js'; +import { SessionManager, tokenStatsManager } from '../../session/index.js'; import { getAgentManager } from '../../agent/manager.js'; import { loadVisionConfig } from '../../utils/config.js'; import { loadDescription } from '../load_description.js'; @@ -317,6 +317,22 @@ async function executeTask(params: TaskParams): Promise { ]; await sessionManager.saveChildSession(childSession); + // 更新子会话的 Token 统计 + const session = sessionManager.getSession(); + if (result.usage && session) { + await tokenStatsManager.updateSessionStats( + session.projectId, + childSession.id, + result.usage + ); + // 合并子会话统计到父会话 + await tokenStatsManager.mergeChildSessionStats( + session.projectId, + parentSessionId, + childSession.id + ); + } + if (result.success) { return { success: true, diff --git a/packages/core/src/types/index.ts b/packages/core/src/types/index.ts index 4c44976..df2bf35 100644 --- a/packages/core/src/types/index.ts +++ b/packages/core/src/types/index.ts @@ -96,12 +96,30 @@ export interface ConversationContext { workingDirectory: string; } +/** + * Token 使用统计信息 + */ +export interface TokenUsageInfo { + /** 输入 tokens（prompt） */ + promptTokens: number; + /** 输出 tokens（completion） */ + completionTokens: number; + /** 总 tokens */ + totalTokens: number; + /** 缓存读取的输入 tokens（Anthropic API） */ + cacheReadInputTokens?: number; + /** 缓存创建的输入 tokens（Anthropic API） */ + cacheCreationInputTokens?: number; +} + // Chat 返回结果（包含完整的消息链） export interface ChatResult { /** 最终文本响应 */ text: string; /** 完整的响应消息链（包含 tool-call 和 tool-result） */ messages: unknown[]; + /** Token 使用统计（从 AI SDK response.usage 提取） */ + usage?: TokenUsageInfo; } // 将自定义 Tool 转换为 Vercel AI SDK 的 zod schema diff --git a/packages/core/tests/unit/session/token-stats-manager.test.ts b/packages/core/tests/unit/session/token-stats-manager.test.ts new file mode 100644 index 0000000..3916b0c --- /dev/null +++ b/packages/core/tests/unit/session/token-stats-manager.test.ts @@ -0,0 +1,372 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest'; + +// Mock storage functions before imports +vi.mock('../../../src/session/storage/session.js', () => ({ + get: vi.fn(), + update: vi.fn(), + updateStats: vi.fn(), + list: vi.fn(), + listByProject: vi.fn(), +})); + +vi.mock('../../../src/session/storage/index.js', () => ({ + read: vi.fn(), + update: vi.fn(), + StorageNotFoundError: class StorageNotFoundError extends Error { + constructor(message: string) { + super(message); + this.name = 'StorageNotFoundError'; + } + }, +})); + +import * as SessionStorage from '../../../src/session/storage/session.js'; +import { TokenStatsManager } from '../../../src/session/token-stats-manager.js'; +import type { TokenUsageInfo } from '../../../src/types/index.js'; + +describe('TokenStatsManager', () => { + let manager: TokenStatsManager; + + beforeEach(() => { + vi.clearAllMocks(); + manager = new TokenStatsManager(); + }); + + describe('updateSessionStats', () => { + it('会话不存在时不更新', async () => { + const usage: TokenUsageInfo = { + promptTokens: 1000, + completionTokens: 500, + totalTokens: 1500, + }; + + vi.mocked(SessionStorage.get).mockResolvedValue(null); + + await manager.updateSessionStats('project-1', 'session-1', usage); + + expect(SessionStorage.get).toHaveBeenCalledWith('project-1', 'session-1'); + expect(SessionStorage.updateStats).not.toHaveBeenCalled(); + }); + + it('更新会话统计 - 已存在的会话', async () => { + const usage: TokenUsageInfo = { + promptTokens: 1000, + completionTokens: 500, + totalTokens: 1500, + }; + + vi.mocked(SessionStorage.get).mockResolvedValue({ + id: 'session-1', + projectId: 'project-1', + createdAt: Date.now(), + updatedAt: Date.now(), + workdir: '/test', + stats: { + messageCount: 5, + inputTokens: 2000, + outputTokens: 1000, + totalTokens: 3000, + updatedAt: Date.now() - 1000, + }, + }); + + await manager.updateSessionStats('project-1', 'session-1', usage); + + expect(SessionStorage.updateStats).toHaveBeenCalledWith( + 'project-1', + 'session-1', + expect.objectContaining({ + messageCount: 6, // 5 + 1 + inputTokens: 3000, // 2000 + 1000 + outputTokens: 1500, // 1000 + 500 + totalTokens: 4500, // 3000 + 1500 + }) + ); + }); + + it('处理缓存 token', async () => { + const usage: TokenUsageInfo = { + promptTokens: 1000, + completionTokens: 500, + totalTokens: 1500, + cacheReadInputTokens: 200, + cacheCreationInputTokens: 100, + }; + + vi.mocked(SessionStorage.get).mockResolvedValue({ + id: 'session-1', + projectId: 'project-1', + createdAt: Date.now(), + updatedAt: Date.now(), + workdir: '/test', + stats: { + messageCount: 0, + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + updatedAt: Date.now(), + }, + }); + + await manager.updateSessionStats('project-1', 'session-1', usage); + + expect(SessionStorage.updateStats).toHaveBeenCalledWith( + 'project-1', + 'session-1', + expect.objectContaining({ + cacheReadTokens: 200, + cacheWriteTokens: 100, + }) + ); + }); + }); + + describe('mergeChildSessionStats', () => { + it('合并子会话统计到父会话', async () => { + vi.mocked(SessionStorage.get).mockImplementation(async (projectId: string, sessionId: string) => { + if (sessionId === 'child-session') { + return { + id: 'child-session', + projectId, + createdAt: Date.now(), + updatedAt: Date.now(), + workdir: '/test', + stats: { + messageCount: 3, + inputTokens: 500, + outputTokens: 300, + totalTokens: 800, + updatedAt: Date.now(), + }, + }; + } + if (sessionId === 'parent-session') { + return { + id: 'parent-session', + projectId, + createdAt: Date.now(), + updatedAt: Date.now(), + workdir: '/test', + stats: { + messageCount: 10, + inputTokens: 2000, + outputTokens: 1000, + totalTokens: 3000, + updatedAt: Date.now(), + }, + }; + } + return null; + }); + + await manager.mergeChildSessionStats('project-1', 'parent-session', 'child-session'); + + expect(SessionStorage.updateStats).toHaveBeenCalledWith( + 'project-1', + 'parent-session', + expect.objectContaining({ + childSessionsTokens: expect.objectContaining({ + inputTokens: 500, + outputTokens: 300, + totalTokens: 800, + }), + }) + ); + }); + + it('累加子会话统计', async () => { + vi.mocked(SessionStorage.get).mockImplementation(async (projectId: string, sessionId: string) => { + if (sessionId === 'child-session-2') { + return { + id: 'child-session-2', + projectId, + createdAt: Date.now(), + updatedAt: Date.now(), + workdir: '/test', + stats: { + messageCount: 2, + inputTokens: 300, + outputTokens: 200, + totalTokens: 500, + updatedAt: Date.now(), + }, + }; + } + if (sessionId === 'parent-session') { + return { + id: 'parent-session', + projectId, + createdAt: Date.now(), + updatedAt: Date.now(), + workdir: '/test', + stats: { + messageCount: 10, + inputTokens: 2000, + outputTokens: 1000, + totalTokens: 3000, + childSessionsTokens: { + inputTokens: 500, + outputTokens: 300, + totalTokens: 800, + }, + updatedAt: Date.now(), + }, + }; + } + return null; + }); + + await manager.mergeChildSessionStats('project-1', 'parent-session', 'child-session-2'); + + expect(SessionStorage.updateStats).toHaveBeenCalledWith( + 'project-1', + 'parent-session', + expect.objectContaining({ + childSessionsTokens: expect.objectContaining({ + inputTokens: 800, // 500 + 300 + outputTokens: 500, // 300 + 200 + totalTokens: 1300, // 800 + 500 + }), + }) + ); + }); + }); + + describe('getSessionStats', () => { + it('返回完整的会话统计', async () => { + const now = Date.now(); + vi.mocked(SessionStorage.get).mockResolvedValue({ + id: 'session-1', + projectId: 'project-1', + createdAt: now, + updatedAt: now, + workdir: '/test', + stats: { + messageCount: 10, + inputTokens: 5000, + outputTokens: 2500, + totalTokens: 7500, + cacheReadTokens: 1000, + cacheWriteTokens: 500, + childSessionsTokens: { + inputTokens: 1000, + outputTokens: 500, + totalTokens: 1500, + }, + updatedAt: now, + }, + }); + + const stats = await manager.getSessionStats('project-1', 'session-1'); + + expect(stats).toMatchObject({ + self: { + inputTokens: 5000, + outputTokens: 2500, + totalTokens: 7500, + cacheReadTokens: 1000, + cacheWriteTokens: 500, + }, + children: { + inputTokens: 1000, + outputTokens: 500, + totalTokens: 1500, + }, + total: { + inputTokens: 6000, // 5000 + 1000 + outputTokens: 3000, // 2500 + 500 + totalTokens: 9000, // 7500 + 1500 + }, + messageCount: 10, + }); + }); + + it('会话不存在时返回 null', async () => { + vi.mocked(SessionStorage.get).mockResolvedValue(null); + + const stats = await manager.getSessionStats('project-1', 'nonexistent'); + + expect(stats).toBeNull(); + }); + + it('没有子会话统计时返回空的 children', async () => { + vi.mocked(SessionStorage.get).mockResolvedValue({ + id: 'session-1', + projectId: 'project-1', + createdAt: Date.now(), + updatedAt: Date.now(), + workdir: '/test', + stats: { + messageCount: 5, + inputTokens: 1000, + outputTokens: 500, + totalTokens: 1500, + updatedAt: Date.now(), + }, + }); + + const stats = await manager.getSessionStats('project-1', 'session-1'); + + expect(stats?.children).toEqual({ + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + }); + expect(stats?.total).toEqual({ + inputTokens: 1000, + outputTokens: 500, + totalTokens: 1500, + }); + }); + }); + + describe('formatTokens', () => { + it('格式化小数字', () => { + expect(manager.formatTokens(100)).toBe('100'); + expect(manager.formatTokens(999)).toBe('999'); + }); + + it('格式化千级数字', () => { + expect(manager.formatTokens(1000)).toBe('1.0k'); + expect(manager.formatTokens(1500)).toBe('1.5k'); + expect(manager.formatTokens(9999)).toBe('10.0k'); + }); + + it('格式化百万级数字', () => { + expect(manager.formatTokens(1000000)).toBe('1.00M'); + expect(manager.formatTokens(1500000)).toBe('1.50M'); + }); + }); + + describe('formatSessionStats', () => { + it('格式化会话统计为人类可读字符串', () => { + const stats = { + self: { + inputTokens: 5000, + outputTokens: 2500, + totalTokens: 7500, + }, + children: { + inputTokens: 1000, + outputTokens: 500, + totalTokens: 1500, + }, + total: { + inputTokens: 6000, + outputTokens: 3000, + totalTokens: 9000, + }, + messageCount: 10, + updatedAt: Date.now(), + }; + + const formatted = manager.formatSessionStats(stats); + + expect(formatted).toContain('6.0k'); // total input + expect(formatted).toContain('3.0k'); // total output + expect(formatted).toContain('9.0k'); // total + expect(formatted).toContain('1.5k'); // children + expect(formatted).toContain('10'); // message count + }); + }); +}); diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts index 3cd5734..f3899d7 100644 --- a/packages/server/src/index.ts +++ b/packages/server/src/index.ts @@ -9,7 +9,7 @@ import { cors } from 'hono/cors'; import { logger } from 'hono/logger'; import { createBunWebSocket } from 'hono/bun'; -import { sessionsRouter, toolsRouter, configRouter, filesRouter, commandsRouter, mcpRouter, hooksRouter, agentsRouter, checkpointsRouter, providersRouter, servicesRouter, contextRouter, lspRouter, systemCommandsRouter } from './routes/index.js'; +import { sessionsRouter, toolsRouter, configRouter, filesRouter, commandsRouter, mcpRouter, hooksRouter, agentsRouter, checkpointsRouter, providersRouter, servicesRouter, contextRouter, lspRouter, systemCommandsRouter, statsRouter } from './routes/index.js'; import { handleWebSocket, handleWebSocketMessage, @@ -91,6 +91,7 @@ api.route('/providers', providersRouter); api.route('/services', servicesRouter); api.route('/lsp', lspRouter); api.route('/system-commands', systemCommandsRouter); +api.route('/stats', statsRouter); // 上下文压缩相关（挂载到根路径，内部路由包含 /sessions/:id/context） api.route('/', contextRouter); diff --git a/packages/server/src/routes/index.ts b/packages/server/src/routes/index.ts index 620229f..191a50d 100644 --- a/packages/server/src/routes/index.ts +++ b/packages/server/src/routes/index.ts @@ -18,3 +18,4 @@ export { servicesRouter } from './services.js'; export { contextRouter } from './context.js'; export { lspRouter } from './lsp.js'; export { systemCommandsRouter } from './system-commands.js'; +export { statsRouter } from './stats.js'; diff --git a/packages/server/src/routes/stats.ts b/packages/server/src/routes/stats.ts new file mode 100644 index 0000000..5d94c38 --- /dev/null +++ b/packages/server/src/routes/stats.ts @@ -0,0 +1,216 @@ +/** + * Stats API Routes + * + * Token 消耗统计相关的 REST API + */ + +import { Hono } from 'hono'; +import { tokenStatsManager } from '@ai-assistant/core'; +import { getSessionManager } from '../session/manager.js'; + +export const statsRouter = new Hono(); + +const sessionManager = getSessionManager(); + +/** + * GET /stats/sessions/:id - 获取会话 Token 统计 + */ +statsRouter.get('/sessions/:id', async (c) => { + const id = c.req.param('id'); + + if (!sessionManager.exists(id)) { + return c.json( + { + success: false, + error: 'Session not found', + }, + 404 + ); + } + + try { + // 从 session manager 获取 projectId + const projectId = sessionManager.getProjectId(id); + + const stats = await tokenStatsManager.getSessionStats(projectId, id); + + if (!stats) { + // 如果没有统计，返回空数据 + return c.json({ + success: true, + data: { + sessionId: id, + self: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + }, + children: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + }, + total: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + }, + messageCount: 0, + }, + }); + } + + return c.json({ + success: true, + data: stats, + }); + } catch (error) { + console.error('[Stats] Failed to get session stats:', error); + return c.json( + { + success: false, + error: error instanceof Error ? error.message : 'Failed to get stats', + }, + 500 + ); + } +}); + +/** + * GET /stats/projects/:id - 获取项目 Token 统计 + */ +statsRouter.get('/projects/:id', async (c) => { + const id = c.req.param('id'); + + try { + // forceRefresh=true 强制刷新统计 + const refresh = c.req.query('refresh') === 'true'; + const stats = await tokenStatsManager.getProjectStats(id, refresh); + + if (!stats) { + return c.json({ + success: true, + data: { + projectId: id, + totalInputTokens: 0, + totalOutputTokens: 0, + totalTokens: 0, + sessionCount: 0, + }, + }); + } + + return c.json({ + success: true, + data: { + projectId: id, + ...stats, + }, + }); + } catch (error) { + console.error('[Stats] Failed to get project stats:', error); + return c.json( + { + success: false, + error: error instanceof Error ? error.message : 'Failed to get stats', + }, + 500 + ); + } +}); + +/** + * GET /stats/summary - 获取当前会话/项目统计摘要 + * + * 返回当前活动会话的统计信息（如果有） + */ +statsRouter.get('/summary', async (c) => { + try { + const sessions = sessionManager.list(); + const currentSession = sessions.find((s) => s.status === 'busy') || sessions[0]; + + if (!currentSession) { + return c.json({ + success: true, + data: { + hasActiveSession: false, + session: null, + project: null, + }, + }); + } + + // 从 session manager 获取 projectId + const projectId = sessionManager.getProjectId(currentSession.id); + + const [sessionStats, projectStats] = await Promise.all([ + tokenStatsManager.getSessionStats(projectId, currentSession.id), + tokenStatsManager.getProjectStats(projectId), + ]); + + return c.json({ + success: true, + data: { + hasActiveSession: true, + sessionId: currentSession.id, + projectId, + session: sessionStats || { + sessionId: currentSession.id, + self: { inputTokens: 0, outputTokens: 0, totalTokens: 0 }, + children: { inputTokens: 0, outputTokens: 0, totalTokens: 0 }, + total: { inputTokens: 0, outputTokens: 0, totalTokens: 0 }, + messageCount: 0, + }, + project: projectStats + ? { + projectId, + ...projectStats, + } + : { + projectId, + totalInputTokens: 0, + totalOutputTokens: 0, + totalTokens: 0, + sessionCount: 0, + }, + }, + }); + } catch (error) { + console.error('[Stats] Failed to get summary:', error); + return c.json( + { + success: false, + error: error instanceof Error ? error.message : 'Failed to get summary', + }, + 500 + ); + } +}); + +/** + * POST /stats/projects/:id/refresh - 强制刷新项目统计 + */ +statsRouter.post('/projects/:id/refresh', async (c) => { + const id = c.req.param('id'); + + try { + const stats = await tokenStatsManager.updateProjectStats(id); + + return c.json({ + success: true, + data: { + projectId: id, + ...stats, + }, + }); + } catch (error) { + console.error('[Stats] Failed to refresh project stats:', error); + return c.json( + { + success: false, + error: error instanceof Error ? error.message : 'Failed to refresh stats', + }, + 500 + ); + } +});