fix: cache warm transcript weights

This commit is contained in:
Xipong
2026-09-09 01:08:35 +03:00
committed by Teknium
parent 5a2ccffd41
commit 3bfd7b653b
2 changed files with 51 additions and 2 deletions
@@ -1,4 +1,4 @@
import { beforeEach, describe, expect, it } from 'vitest'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import type { ClientSessionState } from '@/app/types'
import type { ChatMessage } from '@/lib/chat-messages'
@@ -64,6 +64,41 @@ describe('SessionStateCache', () => {
expect(owners.get('stored-a')).toBe('runtime-new-owner')
})
it('does not remeasure unchanged warm transcripts on repeated stream flushes', () => {
const stringify = vi.spyOn(JSON, 'stringify')
const cache = new SessionStateCache(
{ isReferenced: () => false, onEvict: () => undefined },
{ maxBytes: Number.POSITIVE_INFINITY, maxCount: 24 }
)
for (let index = 0; index < 24; index += 1) {
cache.set(`runtime-${index}`, settled(`stored-${index}`, 'x'.repeat(4096)))
}
for (let flush = 0; flush < 10; flush += 1) {
cache.prune()
}
expect(stringify).toHaveBeenCalledTimes(24)
stringify.mockRestore()
})
it('never measures a pinned active transcript while streaming', () => {
const cache = new SessionStateCache({ isReferenced: () => true, onEvict: () => undefined })
const state = settled('active', 'stream')
const stringify = vi.spyOn(JSON, 'stringify')
try {
for (let flush = 0; flush < 20; flush += 1) {
cache.set('active', { ...state, messages: [...state.messages] })
cache.prune()
}
expect(stringify).not.toHaveBeenCalled()
} finally {
stringify.mockRestore()
}
})
it('uses transcript bytes as well as count', () => {
const evicted: string[] = []
@@ -45,6 +45,7 @@ export class SessionStateCache extends Map<string, ClientSessionState> {
readonly #maxBytes: number
readonly #maxCount: number
readonly #recency = new Map<string, number>()
readonly #transcriptWeights = new WeakMap<ClientSessionState['messages'], number>()
#clock = 0
constructor(callbacks: SessionStateCacheCallbacks, limits: SessionStateCacheLimits = {}) {
@@ -91,7 +92,7 @@ export class SessionStateCache extends Map<string, ClientSessionState> {
continue
}
const weight = transcriptBytes(state)
const weight = this.#weight(state)
candidates.push({ bytes: weight, runtimeId, state, touched: this.#recency.get(runtimeId) ?? 0 })
bytes += weight
}
@@ -145,6 +146,19 @@ export class SessionStateCache extends Map<string, ClientSessionState> {
)
}
#weight(state: ClientSessionState): number {
const cached = this.#transcriptWeights.get(state.messages)
if (cached !== undefined) {
return cached
}
const weight = transcriptBytes(state)
this.#transcriptWeights.set(state.messages, weight)
return weight
}
#touch(runtimeId: string): void {
this.#clock += 1
this.#recency.set(runtimeId, this.#clock)