Something went wrong. Try again.
source dump of claude code forked from oppi.li/claude-code
Something went wrong. Try again.
17 kB · 495 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496/** * Session Memory automatically maintains a markdown file with notes about the current conversation. * It runs periodically in the background using a forked subagent to extract key information * without interrupting the main conversation flow. */
import { writeFile } from 'fs/promises'import memoize from 'lodash-es/memoize.js'import { getIsRemoteMode } from '../../bootstrap/state.js'import { getSystemPrompt } from '../../constants/prompts.js'import { getSystemContext, getUserContext } from '../../context.js'import type { CanUseToolFn } from '../../hooks/useCanUseTool.js'import type { Tool, ToolUseContext } from '../../Tool.js'import { FILE_EDIT_TOOL_NAME } from '../../tools/FileEditTool/constants.js'import { FileReadTool, type Output as FileReadToolOutput,} from '../../tools/FileReadTool/FileReadTool.js'import type { Message } from '../../types/message.js'import { count } from '../../utils/array.js'import { createCacheSafeParams, createSubagentContext, runForkedAgent,} from '../../utils/forkedAgent.js'import { getFsImplementation } from '../../utils/fsOperations.js'import { type REPLHookContext, registerPostSamplingHook,} from '../../utils/hooks/postSamplingHooks.js'import { createUserMessage, hasToolCallsInLastAssistantTurn,} from '../../utils/messages.js'import { getSessionMemoryDir, getSessionMemoryPath,} from '../../utils/permissions/filesystem.js'import { sequential } from '../../utils/sequential.js'import { asSystemPrompt } from '../../utils/systemPromptType.js'import { getTokenUsage, tokenCountWithEstimation } from '../../utils/tokens.js'import { logEvent } from '../analytics/index.js'import { isAutoCompactEnabled } from '../compact/autoCompact.js'import { buildSessionMemoryUpdatePrompt, loadSessionMemoryTemplate,} from './prompts.js'import { DEFAULT_SESSION_MEMORY_CONFIG, getSessionMemoryConfig, getToolCallsBetweenUpdates, hasMetInitializationThreshold, hasMetUpdateThreshold, isSessionMemoryInitialized, markExtractionCompleted, markExtractionStarted, markSessionMemoryInitialized, recordExtractionTokenCount, type SessionMemoryConfig, setLastSummarizedMessageId, setSessionMemoryConfig,} from './sessionMemoryUtils.js'
// ============================================================================// Feature Gate and Config (Cached - Non-blocking)// ============================================================================// These functions return cached values from disk immediately without blocking// on GrowthBook initialization. Values may be stale but are updated in background.
import { errorMessage, getErrnoCode } from '../../utils/errors.js'import { getDynamicConfig_CACHED_MAY_BE_STALE, getFeatureValue_CACHED_MAY_BE_STALE,} from '../analytics/growthbook.js'
/** * Check if session memory feature is enabled. * Uses cached gate value - returns immediately without blocking. */function isSessionMemoryGateEnabled(): boolean { return getFeatureValue_CACHED_MAY_BE_STALE('tengu_session_memory', false)}
/** * Get session memory config from cache. * Returns immediately without blocking - value may be stale. */function getSessionMemoryRemoteConfig(): Partial<SessionMemoryConfig> { return getDynamicConfig_CACHED_MAY_BE_STALE<Partial<SessionMemoryConfig>>( 'tengu_sm_config', {}, )}
// ============================================================================// Module State// ============================================================================
let lastMemoryMessageUuid: string | undefined
/** * Reset the last memory message UUID (for testing) */export function resetLastMemoryMessageUuid(): void { lastMemoryMessageUuid = undefined}
function countToolCallsSince( messages: Message[], sinceUuid: string | undefined,): number { let toolCallCount = 0 let foundStart = sinceUuid === null || sinceUuid === undefined
for (const message of messages) { if (!foundStart) { if (message.uuid === sinceUuid) { foundStart = true } continue }
if (message.type === 'assistant') { const content = message.message.content if (Array.isArray(content)) { toolCallCount += count(content, block => block.type === 'tool_use') } } }
return toolCallCount}
export function shouldExtractMemory(messages: Message[]): boolean { // Check if we've met the initialization threshold // Uses total context window tokens (same as autocompact) for consistent behavior const currentTokenCount = tokenCountWithEstimation(messages) if (!isSessionMemoryInitialized()) { if (!hasMetInitializationThreshold(currentTokenCount)) { return false } markSessionMemoryInitialized() }
// Check if we've met the minimum tokens between updates threshold // Uses context window growth since last extraction (same metric as init threshold) const hasMetTokenThreshold = hasMetUpdateThreshold(currentTokenCount)
// Check if we've met the tool calls threshold const toolCallsSinceLastUpdate = countToolCallsSince( messages, lastMemoryMessageUuid, ) const hasMetToolCallThreshold = toolCallsSinceLastUpdate >= getToolCallsBetweenUpdates()
// Check if the last assistant turn has no tool calls (safe to extract) const hasToolCallsInLastTurn = hasToolCallsInLastAssistantTurn(messages)
// Trigger extraction when: // 1. Both thresholds are met (tokens AND tool calls), OR // 2. No tool calls in last turn AND token threshold is met // (to ensure we extract at natural conversation breaks) // // IMPORTANT: The token threshold (minimumTokensBetweenUpdate) is ALWAYS required. // Even if the tool call threshold is met, extraction won't happen until the // token threshold is also satisfied. This prevents excessive extractions. const shouldExtract = (hasMetTokenThreshold && hasMetToolCallThreshold) || (hasMetTokenThreshold && !hasToolCallsInLastTurn)
if (shouldExtract) { const lastMessage = messages[messages.length - 1] if (lastMessage?.uuid) { lastMemoryMessageUuid = lastMessage.uuid } return true }
return false}
async function setupSessionMemoryFile( toolUseContext: ToolUseContext,): Promise<{ memoryPath: string; currentMemory: string }> { const fs = getFsImplementation()
// Set up directory and file const sessionMemoryDir = getSessionMemoryDir() await fs.mkdir(sessionMemoryDir, { mode: 0o700 })
const memoryPath = getSessionMemoryPath()
// Create the memory file if it doesn't exist (wx = O_CREAT|O_EXCL) try { await writeFile(memoryPath, '', { encoding: 'utf-8', mode: 0o600, flag: 'wx', }) // Only load template if file was just created const template = await loadSessionMemoryTemplate() await writeFile(memoryPath, template, { encoding: 'utf-8', mode: 0o600, }) } catch (e: unknown) { const code = getErrnoCode(e) if (code !== 'EEXIST') { throw e } }
// Drop any cached entry so FileReadTool's dedup doesn't return a // file_unchanged stub — we need the actual content. The Read repopulates it. toolUseContext.readFileState.delete(memoryPath) const result = await FileReadTool.call( { file_path: memoryPath }, toolUseContext, ) let currentMemory = ''
const output = result.data as FileReadToolOutput if (output.type === 'text') { currentMemory = output.file.content }
logEvent('tengu_session_memory_file_read', { content_length: currentMemory.length, })
return { memoryPath, currentMemory }}
/** * Initialize session memory config from remote config (lazy initialization). * Memoized - only runs once per session, subsequent calls return immediately. * Uses cached config values - non-blocking. */const initSessionMemoryConfigIfNeeded = memoize((): void => { // Load config from cache (non-blocking, may be stale) const remoteConfig = getSessionMemoryRemoteConfig()
// Only use remote values if they are explicitly set (non-zero positive numbers) // This ensures sensible defaults aren't overridden by zero values const config: SessionMemoryConfig = { minimumMessageTokensToInit: remoteConfig.minimumMessageTokensToInit && remoteConfig.minimumMessageTokensToInit > 0 ? remoteConfig.minimumMessageTokensToInit : DEFAULT_SESSION_MEMORY_CONFIG.minimumMessageTokensToInit, minimumTokensBetweenUpdate: remoteConfig.minimumTokensBetweenUpdate && remoteConfig.minimumTokensBetweenUpdate > 0 ? remoteConfig.minimumTokensBetweenUpdate : DEFAULT_SESSION_MEMORY_CONFIG.minimumTokensBetweenUpdate, toolCallsBetweenUpdates: remoteConfig.toolCallsBetweenUpdates && remoteConfig.toolCallsBetweenUpdates > 0 ? remoteConfig.toolCallsBetweenUpdates : DEFAULT_SESSION_MEMORY_CONFIG.toolCallsBetweenUpdates, } setSessionMemoryConfig(config)})
/** * Session memory post-sampling hook that extracts and updates session notes */// Track if we've logged the gate check failure this session (to avoid spam)let hasLoggedGateFailure = false
const extractSessionMemory = sequential(async function ( context: REPLHookContext,): Promise<void> { const { messages, toolUseContext, querySource } = context
// Only run session memory on main REPL thread if (querySource !== 'repl_main_thread') { // Don't log this - it's expected for subagents, teammates, etc. return }
// Check gate lazily when hook runs (cached, non-blocking) if (!isSessionMemoryGateEnabled()) { // Log gate failure once per session (ant-only) if (process.env.USER_TYPE === 'ant' && !hasLoggedGateFailure) { hasLoggedGateFailure = true logEvent('tengu_session_memory_gate_disabled', {}) } return }
// Initialize config from remote (lazy, only once) initSessionMemoryConfigIfNeeded()
if (!shouldExtractMemory(messages)) { return }
markExtractionStarted()
// Create isolated context for setup to avoid polluting parent's cache const setupContext = createSubagentContext(toolUseContext)
// Set up file system and read current state with isolated context const { memoryPath, currentMemory } = await setupSessionMemoryFile(setupContext)
// Create extraction message const userPrompt = await buildSessionMemoryUpdatePrompt( currentMemory, memoryPath, )
// Run session memory extraction using runForkedAgent for prompt caching // runForkedAgent creates an isolated context to prevent mutation of parent state // Pass setupContext.readFileState so the forked agent can edit the memory file await runForkedAgent({ promptMessages: [createUserMessage({ content: userPrompt })], cacheSafeParams: createCacheSafeParams(context), canUseTool: createMemoryFileCanUseTool(memoryPath), querySource: 'session_memory', forkLabel: 'session_memory', overrides: { readFileState: setupContext.readFileState }, })
// Log extraction event for tracking frequency // Use the token usage from the last message in the conversation const lastMessage = messages[messages.length - 1] const usage = lastMessage ? getTokenUsage(lastMessage) : undefined const config = getSessionMemoryConfig() logEvent('tengu_session_memory_extraction', { input_tokens: usage?.input_tokens, output_tokens: usage?.output_tokens, cache_read_input_tokens: usage?.cache_read_input_tokens ?? undefined, cache_creation_input_tokens: usage?.cache_creation_input_tokens ?? undefined, config_min_message_tokens_to_init: config.minimumMessageTokensToInit, config_min_tokens_between_update: config.minimumTokensBetweenUpdate, config_tool_calls_between_updates: config.toolCallsBetweenUpdates, })
// Record the context size at extraction for tracking minimumTokensBetweenUpdate recordExtractionTokenCount(tokenCountWithEstimation(messages))
// Update lastSummarizedMessageId after successful completion updateLastSummarizedMessageIdIfSafe(messages)
markExtractionCompleted()})
/** * Initialize session memory by registering the post-sampling hook. * This is synchronous to avoid race conditions during startup. * The gate check and config loading happen lazily when the hook runs. */export function initSessionMemory(): void { if (getIsRemoteMode()) return // Session memory is used for compaction, so respect auto-compact settings const autoCompactEnabled = isAutoCompactEnabled()
// Log initialization state (ant-only to avoid noise in external logs) if (process.env.USER_TYPE === 'ant') { logEvent('tengu_session_memory_init', { auto_compact_enabled: autoCompactEnabled, }) }
if (!autoCompactEnabled) { return }
// Register hook unconditionally - gate check happens lazily when hook runs registerPostSamplingHook(extractSessionMemory)}
export type ManualExtractionResult = { success: boolean memoryPath?: string error?: string}
/** * Manually trigger session memory extraction, bypassing threshold checks. * Used by the /summary command. */export async function manuallyExtractSessionMemory( messages: Message[], toolUseContext: ToolUseContext,): Promise<ManualExtractionResult> { if (messages.length === 0) { return { success: false, error: 'No messages to summarize' } } markExtractionStarted()
try { // Create isolated context for setup to avoid polluting parent's cache const setupContext = createSubagentContext(toolUseContext)
// Set up file system and read current state with isolated context const { memoryPath, currentMemory } = await setupSessionMemoryFile(setupContext)
// Create extraction message const userPrompt = await buildSessionMemoryUpdatePrompt( currentMemory, memoryPath, )
// Get system prompt for cache-safe params const { tools, mainLoopModel } = toolUseContext.options const [rawSystemPrompt, userContext, systemContext] = await Promise.all([ getSystemPrompt(tools, mainLoopModel), getUserContext(), getSystemContext(), ]) const systemPrompt = asSystemPrompt(rawSystemPrompt)
// Run session memory extraction using runForkedAgent await runForkedAgent({ promptMessages: [createUserMessage({ content: userPrompt })], cacheSafeParams: { systemPrompt, userContext, systemContext, toolUseContext: setupContext, forkContextMessages: messages, }, canUseTool: createMemoryFileCanUseTool(memoryPath), querySource: 'session_memory', forkLabel: 'session_memory_manual', overrides: { readFileState: setupContext.readFileState }, })
// Log manual extraction event logEvent('tengu_session_memory_manual_extraction', {})
// Record the context size at extraction for tracking minimumTokensBetweenUpdate recordExtractionTokenCount(tokenCountWithEstimation(messages))
// Update lastSummarizedMessageId after successful completion updateLastSummarizedMessageIdIfSafe(messages)
return { success: true, memoryPath } } catch (error) { return { success: false, error: errorMessage(error), } } finally { markExtractionCompleted() }}
// Helper functions
/** * Creates a canUseTool function that only allows Edit for the exact memory file. */export function createMemoryFileCanUseTool(memoryPath: string): CanUseToolFn { return async (tool: Tool, input: unknown) => { if ( tool.name === FILE_EDIT_TOOL_NAME && typeof input === 'object' && input !== null && 'file_path' in input ) { const filePath = input.file_path if (typeof filePath === 'string' && filePath === memoryPath) { return { behavior: 'allow' as const, updatedInput: input } } } return { behavior: 'deny' as const, message: `only ${FILE_EDIT_TOOL_NAME} on ${memoryPath} is allowed`, decisionReason: { type: 'other' as const, reason: `only ${FILE_EDIT_TOOL_NAME} on ${memoryPath} is allowed`, }, } }}
/** * Updates lastSummarizedMessageId after successful extraction. * Only sets it if the last message doesn't have tool calls (to avoid orphaned tool_results). */function updateLastSummarizedMessageIdIfSafe(messages: Message[]): void { if (!hasToolCallsInLastAssistantTurn(messages)) { const lastMessage = messages[messages.length - 1] if (lastMessage?.uuid) { setLastSummarizedMessageId(lastMessage.uuid) } }}