diff --git a/apps/desktop/src/main/pi/agent-service.ts b/apps/desktop/src/main/pi/agent-service.ts index e8133e4..44eeefe 100644 --- a/apps/desktop/src/main/pi/agent-service.ts +++ b/apps/desktop/src/main/pi/agent-service.ts @@ -467,13 +467,14 @@ export class AgentService { * 类目：消息 / 自定义工具 / 系统工具 / 系统提示词 / 技能 / 其他； * 技能块注入在系统提示词内，从提示词估算中拆出（不重复计数）。 */ private buildContextInfo(session: AgentSession): ContextUsageInfo { + const agg = this.aggregateSession(session) const breakdown: ContextBreakdownEntry[] = [ - { key: 'messages', label: '消息', tokens: 0 }, - { key: 'customTools', label: '自定义工具', tokens: 0 }, - { key: 'builtinTools', label: '系统工具', tokens: 0 }, + { key: 'messages', label: '消息', tokens: agg.breakdown.get('messages') ?? 0 }, + { key: 'customTools', label: '自定义工具', tokens: agg.breakdown.get('customTools') ?? 0 }, + { key: 'builtinTools', label: '系统工具', tokens: agg.breakdown.get('builtinTools') ?? 0 }, { key: 'systemPrompt', label: '系统提示词', tokens: 0 }, { key: 'skills', label: '技能', tokens: 0 }, - { key: 'other', label: '其他', tokens: 0 } + { key: 'other', label: '其他', tokens: agg.breakdown.get('other') ?? 0 } ] const byKey = new Map(breakdown.map((b) => [b.key, b])) const systemTokens = Math.ceil(session.systemPrompt.length / 4) @@ -481,44 +482,12 @@ export class AgentService { byKey.get('systemPrompt')!.tokens = Math.max(0, systemTokens - skillsTokens) byKey.get('skills')!.tokens = skillsTokens - for (const message of session.messages) { - const m = message as { role?: string; toolName?: unknown; content?: unknown } - switch (m.role) { - case 'user': - case 'assistant': - byKey.get('messages')!.tokens += estimateMessageTokens(message) - break - case 'toolResult': { - const key = - typeof m.toolName === 'string' && CUSTOM_TOOL_NAMES.has(m.toolName) ? 'customTools' : 'builtinTools' - byKey.get(key)!.tokens += estimateMessageTokens(message) - break - } - // bash 执行、扩展注入、压缩/分支摘要等注入上下文的消息 - case 'bashExecution': - case 'custom': - case 'compactionSummary': - case 'branchSummary': - byKey.get('other')!.tokens += estimateMessageTokens(message) - break - default: - break - } - } const usage = session.getContextUsage() - let lastOutputTokens: number | undefined - for (let i = session.messages.length - 1; i >= 0; i--) { - const m = session.messages[i] as { role?: string; usage?: { output?: unknown } } - if (m.role === 'assistant' && m.usage && typeof m.usage.output === 'number') { - lastOutputTokens = m.usage.output - break - } - } return { tokens: usage && typeof usage.tokens === 'number' ? usage.tokens : null, contextWindow: usage?.contextWindow ?? 0, percent: usage && typeof usage.percent === 'number' ? usage.percent : null, - ...(lastOutputTokens !== undefined ? { lastOutputTokens } : {}), + ...(agg.lastOutputTokens !== undefined ? { lastOutputTokens: agg.lastOutputTokens } : {}), breakdown } } @@ -538,42 +507,114 @@ export class AgentService { return value } + // ---- 会话统计增量聚合 ---- + // emitStats 每条 assistant 消息落地都会触发；长会话下每次全量遍历估算 + // token 是 O(全部消息)。消息列表正常只追加：以「已聚合前缀长度 + 前缀 + // 末条消息引用」为锚点，每次只聚合新增尾部；压缩/回溯/会话切换/运行时 + // 重建后消息对象引用必变（或长度回退），锚点失配即自动整体重算。 + + private sessionAgg: { + count: number + anchor: unknown + cacheInput: number + cacheRead: number + cacheWrite: number + toolCalls: number + toolErrors: number + /** 消息侧 breakdown（systemPrompt/skills 在外层另算） */ + breakdown: Map + lastOutputTokens: number | undefined + } | null = null + + private aggregateSession(session: AgentSession) { + const messages = session.messages + let agg = this.sessionAgg + if ( + agg === null || + messages.length < agg.count || + (agg.count > 0 && messages[agg.count - 1] !== agg.anchor) + ) { + agg = { + count: 0, + anchor: null, + cacheInput: 0, + cacheRead: 0, + cacheWrite: 0, + toolCalls: 0, + toolErrors: 0, + breakdown: new Map([ + ['messages', 0], + ['customTools', 0], + ['builtinTools', 0], + ['other', 0] + ]), + lastOutputTokens: undefined + } + } + for (let i = agg.count; i < messages.length; i++) { + const m = messages[i] as { + role?: string + toolName?: unknown + content?: unknown + isError?: unknown + usage?: { input?: unknown; output?: unknown; cacheRead?: unknown; cacheWrite?: unknown } + } + const bump = (key: string): void => { + agg!.breakdown.set(key, (agg!.breakdown.get(key) ?? 0) + estimateMessageTokens(messages[i])) + } + switch (m.role) { + case 'assistant': { + bump('messages') + if (m.usage) { + agg.cacheInput += Number(m.usage.input ?? 0) || 0 + agg.cacheRead += Number(m.usage.cacheRead ?? 0) || 0 + agg.cacheWrite += Number(m.usage.cacheWrite ?? 0) || 0 + if (typeof m.usage.output === 'number') agg.lastOutputTokens = m.usage.output + } + if (Array.isArray(m.content)) { + for (const block of m.content as Array<{ type?: unknown }>) { + if (block && typeof block === 'object' && block.type === 'toolCall') agg.toolCalls += 1 + } + } + break + } + case 'user': + bump('messages') + break + case 'toolResult': { + if (m.isError === true) agg.toolErrors += 1 + bump( + typeof m.toolName === 'string' && CUSTOM_TOOL_NAMES.has(m.toolName) ? 'customTools' : 'builtinTools' + ) + break + } + // bash 执行、扩展注入、压缩/分支摘要等注入上下文的消息 + case 'bashExecution': + case 'custom': + case 'compactionSummary': + case 'branchSummary': + bump('other') + break + default: + break + } + } + agg.count = messages.length + agg.anchor = messages.length > 0 ? messages[messages.length - 1] : null + this.sessionAgg = agg + return agg + } + private async emitStats(): Promise { try { const runtime = await this.getRuntime() const session = runtime.session const stats = session.getSessionStats() const usage = session.getContextUsage() + const agg = this.aggregateSession(session) // 缓存命中率：读缓存 / 全部提示 token（input + cacheRead + cacheWrite）， // 按本会话所有 assistant 响应聚合，0 表示会话内尚无任何响应 - let cacheInput = 0 - let cacheRead = 0 - let cacheWrite = 0 - let toolCalls = 0 - let toolErrors = 0 - for (const message of session.messages) { - const m = message as { - role?: string - usage?: { input?: unknown; cacheRead?: unknown; cacheWrite?: unknown } - content?: unknown - isError?: unknown - } - if (m.role === 'assistant') { - if (m.usage) { - cacheInput += Number(m.usage.input ?? 0) || 0 - cacheRead += Number(m.usage.cacheRead ?? 0) || 0 - cacheWrite += Number(m.usage.cacheWrite ?? 0) || 0 - } - if (Array.isArray(m.content)) { - for (const block of m.content as Array<{ type?: unknown }>) { - if (block && typeof block === 'object' && block.type === 'toolCall') toolCalls += 1 - } - } - } else if (m.role === 'toolResult' && m.isError === true) { - toolErrors += 1 - } - } - const denominator = cacheInput + cacheRead + cacheWrite + const denominator = agg.cacheInput + agg.cacheRead + agg.cacheWrite const info: SessionStatsInfo = { tokens: { ...stats.tokens }, cost: stats.cost, @@ -581,15 +622,15 @@ export class AgentService { context: this.buildContextInfo(session), ...(denominator > 0 ? { - cacheHitRate: Math.round((cacheRead / denominator) * 100), - cacheTokens: { read: cacheRead, write: cacheWrite, input: cacheInput } + cacheHitRate: Math.round((agg.cacheRead / denominator) * 100), + cacheTokens: { read: agg.cacheRead, write: agg.cacheWrite, input: agg.cacheInput } } : {}), - ...(toolCalls > 0 + ...(agg.toolCalls > 0 ? { toolStats: { - calls: toolCalls, - successRate: Math.round(((toolCalls - toolErrors) / toolCalls) * 100) + calls: agg.toolCalls, + successRate: Math.round(((agg.toolCalls - agg.toolErrors) / agg.toolCalls) * 100) } } : {})