Merge pull request #1304 from universe-hcy/fix/cache-safe-params-slot

fix: 修复 lastCacheSafeParams 跨会话残留导致 post-turn fork 上下文污染

Author: claude-code-best

biome.json

@@ -9,6 +9,7 @@
     "includes": [
       "**",
       "!!**/dist",
+      "!!**/dist-stable",
       "!!**/.claude/workflows",
       "!!**/*.workflow.mjs"
     ]

src/cli/print.ts

@@ -185,7 +185,7 @@ import {
   logSuggestionSuppressed,
   type PromptVariant,
 } from 'src/services/PromptSuggestion/promptSuggestion.js'
-import { getLastCacheSafeParams } from 'src/utils/forkedAgent.js'
+import { getLastCacheSafeParams } from 'src/utils/cacheSafeParamsSlot.js'
 import { getAccountInformation } from 'src/utils/auth.js'
 import { OAuthService } from 'src/services/oauth/index.js'
 import { installOAuthTokens } from 'src/cli/handlers/auth.js'

src/commands/btw/btw.tsx

@@ -16,7 +16,8 @@ import type { Message } from '../../types/message.js';
 import { createAbortController } from '../../utils/abortController.js';
 import { saveGlobalConfig } from '../../utils/config.js';
 import { errorMessage } from '../../utils/errors.js';
-import { type CacheSafeParams, getLastCacheSafeParams } from '../../utils/forkedAgent.js';
+import { getLastCacheSafeParams } from '../../utils/cacheSafeParamsSlot.js';
+import type { CacheSafeParams } from '../../utils/forkedAgent.js';
 import { getMessagesAfterCompactBoundary } from '../../utils/messages.js';
 import type { ProcessUserInputContext } from '../../utils/processUserInput/processUserInput.js';
 import { runSideQuestion } from '../../utils/sideQuestion.js';

src/commands/clear/conversation.ts

@@ -29,6 +29,7 @@ import {
 import { isLocalShellTask } from '../../tasks/LocalShellTask/guards.js'
 import { asAgentId } from '../../types/ids.js'
 import type { Message } from '../../types/message.js'
+import { saveCacheSafeParams } from '../../utils/cacheSafeParamsSlot.js'
 import { createEmptyAttributionState } from '../../utils/commitAttribution.js'
 import type { FileStateCache } from '../../utils/fileStateCache.js'
 import {
@@ -156,6 +157,12 @@ export async function clearConversation({
   setLastClassifierRequests(null)
   resetCostState()
 
+  // Drop the post-turn CacheSafeParams snapshot: it holds the pre-clear
+  // conversation's full message history. The session-id check in
+  // getLastCacheSafeParams would reject it anyway after regenerateSessionId
+  // below, but clearing here releases the memory immediately.
+  saveCacheSafeParams(null)
+
   setCwd(getOriginalCwd())
   readFileState.clear()
   discoveredSkillNames?.clear()

src/commands/recap/generateRecap.ts

@@ -12,10 +12,8 @@
 
 import { APIUserAbortError } from '@anthropic-ai/sdk'
 import { logForDebugging } from '../../utils/debug.js'
-import {
-  getLastCacheSafeParams,
-  runForkedAgent,
-} from '../../utils/forkedAgent.js'
+import { getLastCacheSafeParams } from '../../utils/cacheSafeParamsSlot.js'
+import { runForkedAgent } from '../../utils/forkedAgent.js'
 import {
   createUserMessage,
   getAssistantMessageText,

src/query/stopHooks.ts

@@ -49,10 +49,8 @@ import type { QuerySource } from '../constants/querySource.js'
 import { executeAutoDream } from '../services/autoDream/autoDream.js'
 import { executePromptSuggestion } from '../services/PromptSuggestion/promptSuggestion.js'
 import { isBareMode, isEnvDefinedFalsy } from '../utils/envUtils.js'
-import {
-  createCacheSafeParams,
-  saveCacheSafeParams,
-} from '../utils/forkedAgent.js'
+import { saveCacheSafeParams } from '../utils/cacheSafeParamsSlot.js'
+import { createCacheSafeParams } from '../utils/forkedAgent.js'
 
 type StopHookResult = {
   blockingErrors: Message[]

src/utils/__tests__/cacheSafeParamsSlot.test.ts

@@ -0,0 +1,91 @@
+/**
+ * The slot holds the last main-session turn's CacheSafeParams (including the
+ * full conversation history in forkContextMessages). It must never survive a
+ * conversation boundary: /clear regenerates the session id and in-process
+ * /resume switches it — in both cases a post-turn fork (/recap, SDK
+ * side_question, SDK promptSuggestion) reading a stale snapshot would send
+ * the previous conversation's history to the API.
+ *
+ * Uses the real bootstrap/state module (no mocks): regenerateSessionId and
+ * switchSession are exactly what /clear and /resume call in production.
+ */
+import { beforeEach, describe, expect, test } from 'bun:test'
+import {
+  getSessionId,
+  regenerateSessionId,
+  switchSession,
+} from 'src/bootstrap/state.js'
+import { asSessionId } from 'src/types/ids.js'
+import {
+  getLastCacheSafeParams,
+  saveCacheSafeParams,
+} from 'src/utils/cacheSafeParamsSlot.js'
+import type { CacheSafeParams } from 'src/utils/forkedAgent.js'
+
+function makeParams(tag: string): CacheSafeParams {
+  // Only identity matters for the slot; the shape is opaque to it.
+  return { forkContextMessages: [tag] } as unknown as CacheSafeParams
+}
+
+beforeEach(() => {
+  saveCacheSafeParams(null)
+})
+
+describe('saveCacheSafeParams / getLastCacheSafeParams', () => {
+  test('returns the saved params within the same session', () => {
+    const params = makeParams('same-session')
+    saveCacheSafeParams(params)
+    expect(getLastCacheSafeParams()).toBe(params)
+    // Repeated reads keep returning it
+    expect(getLastCacheSafeParams()).toBe(params)
+  })
+
+  test('returns null when nothing was saved', () => {
+    expect(getLastCacheSafeParams()).toBeNull()
+  })
+
+  test('saving null clears the slot', () => {
+    saveCacheSafeParams(makeParams('to-clear'))
+    saveCacheSafeParams(null)
+    expect(getLastCacheSafeParams()).toBeNull()
+  })
+
+  test('a later save overwrites an earlier one', () => {
+    saveCacheSafeParams(makeParams('first'))
+    const second = makeParams('second')
+    saveCacheSafeParams(second)
+    expect(getLastCacheSafeParams()).toBe(second)
+  })
+
+  test('snapshot is invalidated by /clear (regenerateSessionId)', () => {
+    saveCacheSafeParams(makeParams('pre-clear'))
+    regenerateSessionId()
+    // Without the session-id check this returned the pre-clear conversation's
+    // history, which /recap and SDK side_question then sent to the API.
+    expect(getLastCacheSafeParams()).toBeNull()
+  })
+
+  test('snapshot is invalidated by in-process /resume (switchSession)', () => {
+    saveCacheSafeParams(makeParams('session-a'))
+    switchSession(asSessionId('00000000-0000-4000-8000-00000000resu'))
+    expect(getLastCacheSafeParams()).toBeNull()
+  })
+
+  test('stale snapshot stays dropped even if the original session id returns', () => {
+    const original = getSessionId()
+    saveCacheSafeParams(makeParams('original-session'))
+    switchSession(asSessionId('00000000-0000-4000-8000-0000000other'))
+    expect(getLastCacheSafeParams()).toBeNull()
+    // Switching back does not resurrect it — the read already released it.
+    switchSession(original)
+    expect(getLastCacheSafeParams()).toBeNull()
+  })
+
+  test('a save after the session switch is valid for the new session', () => {
+    saveCacheSafeParams(makeParams('old'))
+    regenerateSessionId()
+    const fresh = makeParams('new-session')
+    saveCacheSafeParams(fresh)
+    expect(getLastCacheSafeParams()).toBe(fresh)
+  })
+})

src/utils/cacheSafeParamsSlot.ts

@@ -0,0 +1,38 @@
+/**
+ * Process-wide snapshot of the main loop's cache-safe params, written by
+ * handleStopHooks after each main-session turn so post-turn forks
+ * (/recap, SDK side_question, SDK promptSuggestion, /btw) can share the
+ * parent's prompt cache without threading params through every caller.
+ *
+ * The snapshot is only valid for the conversation that wrote it. The session
+ * id is captured at save time and checked at read time: /clear
+ * (regenerateSessionId) and in-process /resume (switchSession) both change
+ * the current session id, which invalidates the snapshot. Without the check,
+ * a post-turn fork started after /clear or /resume would prepend the
+ * previous conversation's full history (forkContextMessages) to its API
+ * request and answer based on a conversation the user no longer sees.
+ *
+ * Kept in its own module (rather than forkedAgent.ts) so the slot can be
+ * unit-tested without forkedAgent's heavy query-loop dependency chain.
+ */
+import { getSessionId } from '../bootstrap/state.js'
+import type { SessionId } from '../types/ids.js'
+import type { CacheSafeParams } from './forkedAgent.js'
+
+let lastCacheSafeParams: CacheSafeParams | null = null
+let savedSessionId: SessionId | null = null
+
+export function saveCacheSafeParams(params: CacheSafeParams | null): void {
+  lastCacheSafeParams = params
+  savedSessionId = params === null ? null : getSessionId()
+}
+
+export function getLastCacheSafeParams(): CacheSafeParams | null {
+  if (lastCacheSafeParams !== null && savedSessionId !== getSessionId()) {
+    // Stale snapshot from a previous conversation — drop it eagerly so the
+    // old message array can be GC'd instead of lingering until the next turn.
+    lastCacheSafeParams = null
+    savedSessionId = null
+  }
+  return lastCacheSafeParams
+}

src/utils/forkedAgent.ts

@@ -71,18 +71,9 @@ export type CacheSafeParams = {
   forkContextMessages: Message[]
 }
 
-// Slot written by handleStopHooks after each turn so post-turn forks
-// (promptSuggestion, postTurnSummary, /btw) can share the main loop's
-// prompt cache without each caller threading params through.
-let lastCacheSafeParams: CacheSafeParams | null = null
-
-export function saveCacheSafeParams(params: CacheSafeParams | null): void {
-  lastCacheSafeParams = params
-}
-
-export function getLastCacheSafeParams(): CacheSafeParams | null {
-  return lastCacheSafeParams
-}
+// The post-turn snapshot slot (saveCacheSafeParams/getLastCacheSafeParams)
+// lives in cacheSafeParamsSlot.ts — it is session-scoped and kept separate
+// so it can be unit-tested without this module's query-loop dependencies.
 
 export type ForkedAgentParams = {
   /** Messages to start the forked query loop with */