sdaduanbilei/cc-haha
1
0
Fork 0
mirror of https://github.com/NanmiCoder/cc-haha synced 2026-07-17 13:13:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(desktop): restore session token usage after reopen (#820)

Browse Source 
Persist transcript usage through session history so reopened desktop tabs can restore header token counts, and resume elapsed timers when reconnected sessions report active status.

Tested: cd desktop && bun run test -- --run src/stores/chatStore.test.ts
Tested: bun test src/server/__tests__/sessions.test.ts
Tested: cd desktop && bun run lint
Tested: bun run check:server
Tested: git diff --check
Not-tested: bun run verify and coverage; scoped local handoff, not PR-ready validation.
Confidence: high
Scope-risk: narrow
This commit is contained in:
程序员阿江(Relakkes) 2026-06-15 23:02:04 +08:00
parent 65e662450c
commit 572eba46f4
5 changed files with 190 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

77
desktop/src/stores/chatStore.test.ts
View File

@ -807,6 +807,49 @@ describe('chatStore history mapping', () => {
})
})
it('restores token usage from transcript history after reopening a session', async () => {
vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({
messages: [
{
id: 'user-1',
type: 'user',
timestamp: '2026-04-06T00:00:00.000Z',
content: 'build the docs',
},
{
id: 'assistant-1',
type: 'assistant',
timestamp: '2026-04-06T00:00:01.000Z',
content: 'done',
usage: { input_tokens: 1200, output_tokens: 80 },
},
{
id: 'assistant-2',
type: 'assistant',
timestamp: '2026-04-06T00:00:02.000Z',
content: [{ type: 'text', text: 'follow-up done' }],
usage: { input_tokens: 3400, output_tokens: 120 },
},
],
})
useChatStore.setState({
sessions: {
[TEST_SESSION_ID]: makeSession({
messages: [],
tokenUsage: { input_tokens: 0, output_tokens: 0 },
}),
},
})
await useChatStore.getState().loadHistory(TEST_SESSION_ID)
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.tokenUsage).toEqual({
input_tokens: 4600,
output_tokens: 200,
})
})
it('uses transcript terminal events to repair stale live goal and background task state', async () => {
vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({
messages: [
@ -2954,6 +2997,40 @@ describe('chatStore history mapping', () => {
expect(updateTabStatusMock).toHaveBeenLastCalledWith(TEST_SESSION_ID, 'running')
})
it('starts an elapsed timer when a reconnected session reports running status', () => {
vi.useFakeTimers()
useChatStore.setState({
sessions: {
[TEST_SESSION_ID]: makeSession({
chatState: 'idle',
elapsedSeconds: 0,
elapsedTimer: null,
}),
},
})
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
type: 'status',
state: 'thinking',
verb: 'Thinking',
})
vi.advanceTimersByTime(2100)
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.elapsedSeconds).toBe(2)
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
type: 'status',
state: 'idle',
})
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.elapsedTimer).toBeNull()
vi.runOnlyPendingTimers()
vi.useRealTimers()
})
it('tracks API retry status until the request finishes', () => {
useChatStore.setState({
sessions: {

50
desktop/src/stores/chatStore.ts
View File

@ -840,6 +840,37 @@ function mergeSlashCommandUpdates(
return [...merged.values()]
}
function readUsageToken(value: unknown): number {
return typeof value === 'number' && Number.isFinite(value) ? value : 0
}
function summarizeTokenUsageFromHistory(messages: MessageEntry[]): TokenUsage | null {
let inputTokens = 0
let outputTokens = 0
let cacheReadTokens = 0
let cacheCreationTokens = 0
for (const message of messages) {
const usage = message.usage
if (!usage) continue
inputTokens += readUsageToken(usage.input_tokens)
outputTokens += readUsageToken(usage.output_tokens)
cacheReadTokens += readUsageToken(usage.cache_read_input_tokens)
cacheCreationTokens += readUsageToken(usage.cache_creation_input_tokens)
}
if (inputTokens === 0 && outputTokens === 0 && cacheReadTokens === 0 && cacheCreationTokens === 0) {
return null
}
return {
input_tokens: inputTokens,
output_tokens: outputTokens,
...(cacheReadTokens > 0 ? { cache_read_tokens: cacheReadTokens } : {}),
...(cacheCreationTokens > 0 ? { cache_creation_tokens: cacheCreationTokens } : {}),
}
}
async function fetchAndMapSessionHistory(sessionId: string) {
const { messages, taskNotifications } = await sessionsApi.getMessages(sessionId)
const uiMessages = mapHistoryMessagesToUiMessages(messages)
@ -855,6 +886,7 @@ async function fetchAndMapSessionHistory(sessionId: string) {
restoredBackgroundTasks: backgroundTaskRecordFromNotifications(Object.values(restoredNotifications)),
lastTodos: extractLastTodoWriteFromHistory(messages),
hasMessagesAfterTaskCompletion: hasUserMessagesAfterTaskCompletion(messages),
tokenUsage: summarizeTokenUsageFromHistory(messages),
}
}
@ -1162,6 +1194,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
restoredBackgroundTasks,
lastTodos,
hasMessagesAfterTaskCompletion,
tokenUsage,
} = await fetchAndMapSessionHistory(sessionId)
set((state) => {
const session = state.sessions[sessionId]
@ -1176,6 +1209,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
s.backgroundAgentTasks ?? {},
restoredBackgroundTasks,
),
tokenUsage: tokenUsage ?? s.tokenUsage,
messages: mergeRestoredHistoryIntoLiveMessages(
mergeBackgroundTaskMessages(s.messages, restoredBackgroundTasks),
uiMessages,
@ -1192,6 +1226,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
s.backgroundAgentTasks ?? {},
restoredBackgroundTasks,
),
tokenUsage: tokenUsage ?? s.tokenUsage,
})) }
})
if (lastTodos && lastTodos.length > 0) {
@ -1235,6 +1270,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
restoredBackgroundTasks,
lastTodos,
hasMessagesAfterTaskCompletion,
tokenUsage,
} = await fetchAndMapSessionHistory(sessionId)
set((state) => {
@ -1249,6 +1285,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
activeGoal,
agentTaskNotifications: restoredNotifications,
backgroundAgentTasks: restoredBackgroundTasks,
tokenUsage: tokenUsage ?? session.tokenUsage,
chatState: 'idle',
activeThinkingId: null,
activeToolUseId: null,
@ -1500,6 +1537,19 @@ export const useChatStore = create<ChatStore>((set, get) => ({
} : pendingText !== session.streamingText ? { streamingText: pendingText } : {}),
}
})
if (msg.state !== 'idle') {
const session = get().sessions[sessionId]
if (session && !session.elapsedTimer) {
const timer = setInterval(() => {
set((st) => ({
sessions: updateSessionIn(st.sessions, sessionId, (sess) => ({
elapsedSeconds: sess.elapsedSeconds + 1,
})),
}))
}, 1000)
update(() => ({ elapsedTimer: timer }))
}
}
if (msg.state === 'idle') {
const session = get().sessions[sessionId]
if (session?.elapsedTimer) {

8
desktop/src/types/session.ts
View File

@ -13,6 +13,13 @@ export type SessionListItem = {
permissionMode?: string
}
export type MessageUsage = {
input_tokens?: number
output_tokens?: number
cache_read_input_tokens?: number
cache_creation_input_tokens?: number
}
export type MessageEntry = {
id: string
type: 'user' | 'assistant' | 'system' | 'tool_use' | 'tool_result'
@ -20,6 +27,7 @@ export type MessageEntry = {
toolUseResult?: unknown
timestamp: string
model?: string
usage?: MessageUsage
parentUuid?: string
parentToolUseId?: string
isSidechain?: boolean

18
src/server/__tests__/sessions.test.ts
View File

@ -225,7 +225,16 @@ function makeUserEntry(content: string, uuid?: string): Record<string, unknown>
}
}
function makeAssistantEntry(content: string, parentUuid?: string): Record<string, unknown> {
function makeAssistantEntry(
content: string,
parentUuid?: string,
usage?: {
input_tokens?: number
output_tokens?: number
cache_read_input_tokens?: number
cache_creation_input_tokens?: number
},
): Record<string, unknown> {
return {
parentUuid: parentUuid || null,
isSidechain: false,
@ -236,6 +245,7 @@ function makeAssistantEntry(content: string, parentUuid?: string): Record<string
type: 'message',
role: 'assistant',
content: [{ type: 'text', text: content }],
...(usage ? { usage } : {}),
},
uuid: crypto.randomUUID(),
timestamp: '2026-01-01T00:02:00.000Z',
@ -2057,7 +2067,7 @@ describe('Sessions API', () => {
await writeSessionFile('-tmp-api-test', sessionId, [
makeSnapshotEntry(),
makeUserEntry('Hello'),
makeAssistantEntry('World'),
makeAssistantEntry('World', undefined, { input_tokens: 1234, output_tokens: 56 }),
makeUserEntry(
'<task-notification>\n<task-id>bg-1</task-id>\n<tool-use-id>toolu_bg</tool-use-id>\n<status>failed</status>\n<summary>Background command failed &amp; stopped</summary>\n<result>Stack trace &amp; failed assertion</result>\n<output-file>C:\\Temp\\bg.output</output-file>\n</task-notification>',
crypto.randomUUID(),
@ -2073,6 +2083,10 @@ describe('Sessions API', () => {
taskNotifications: unknown[]
}
expect(body.messages).toHaveLength(2)
expect(body.messages[1]).toMatchObject({
type: 'assistant',
usage: { input_tokens: 1234, output_tokens: 56 },
})
expect(JSON.stringify(body.messages)).not.toContain('<task-notification>')
expect(body.taskNotifications).toEqual([
{

39
src/server/services/sessionService.ts
View File

@ -94,6 +94,13 @@ export type TrimSessionResult = {
removedMessageIds: string[]
}
export type MessageUsage = {
input_tokens?: number
output_tokens?: number
cache_read_input_tokens?: number
cache_creation_input_tokens?: number
}
export type MessageEntry = {
id: string
type: 'user' | 'assistant' | 'system' | 'tool_use' | 'tool_result'
@ -101,6 +108,7 @@ export type MessageEntry = {
toolUseResult?: unknown
timestamp: string
model?: string
usage?: MessageUsage
parentUuid?: string
parentToolUseId?: string
isSidechain?: boolean
@ -225,6 +233,34 @@ type RawEntry = {
[key: string]: unknown
}
type RawMessageUsage = NonNullable<RawEntry['message']>['usage']
function normalizeMessageUsage(usage: RawMessageUsage): MessageUsage | undefined {
if (!usage) return undefined
const normalized: MessageUsage = {}
if (typeof usage.input_tokens === 'number' && Number.isFinite(usage.input_tokens)) {
normalized.input_tokens = usage.input_tokens
}
if (typeof usage.output_tokens === 'number' && Number.isFinite(usage.output_tokens)) {
normalized.output_tokens = usage.output_tokens
}
if (
typeof usage.cache_read_input_tokens === 'number' &&
Number.isFinite(usage.cache_read_input_tokens)
) {
normalized.cache_read_input_tokens = usage.cache_read_input_tokens
}
if (
typeof usage.cache_creation_input_tokens === 'number' &&
Number.isFinite(usage.cache_creation_input_tokens)
) {
normalized.cache_creation_input_tokens = usage.cache_creation_input_tokens
}
return Object.keys(normalized).length > 0 ? normalized : undefined
}
type PersistedWorktreeSession = {
originalCwd: string
worktreePath: string
@ -700,6 +736,8 @@ export class SessionService {
type = 'system'
}
const usage = normalizeMessageUsage(msg.usage)
return {
id: entry.uuid || crypto.randomUUID(),
type,
@ -707,6 +745,7 @@ export class SessionService {
...(entry.toolUseResult !== undefined ? { toolUseResult: entry.toolUseResult } : {}),
timestamp: entry.timestamp || new Date().toISOString(),
model: msg.model,
...(usage ? { usage } : {}),
parentUuid: entry.parentUuid ?? undefined,
parentToolUseId,
isSidechain: entry.isSidechain,