From 3bfd7b653bb9831270ac30aeb6e72fcf7c43e878 Mon Sep 17 00:00:00 2001
From: Xipong <217837358+Xipong@users.noreply.github.com>
Date: Wed, 9 Sep 2026 01:08:35 +0300
Subject: [PATCH] fix: cache warm transcript weights

---
 .../app/session/session-state-cache.test.ts   | 37 ++++++++++++++++++-
 .../src/app/session/session-state-cache.ts    | 16 +++++++-
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/apps/desktop/src/app/session/session-state-cache.test.ts b/apps/desktop/src/app/session/session-state-cache.test.ts
index ecdeb7f22e..e5f145d078 100644
--- a/apps/desktop/src/app/session/session-state-cache.test.ts
+++ b/apps/desktop/src/app/session/session-state-cache.test.ts
@@ -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[] = []
 
diff --git a/apps/desktop/src/app/session/session-state-cache.ts b/apps/desktop/src/app/session/session-state-cache.ts
index 04208fed72..6047342f44 100644
--- a/apps/desktop/src/app/session/session-state-cache.ts
+++ b/apps/desktop/src/app/session/session-state-cache.ts
@@ -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)