fix: cache warm transcript weights
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user