sessionMemory.ts
1 /** 2 * Session Memory automatically maintains a markdown file with notes about the current conversation. 3 * It runs periodically in the background using a forked subagent to extract key information 4 * without interrupting the main conversation flow. 5 */ 6 7 import { writeFile } from 'fs/promises' 8 import memoize from 'lodash-es/memoize.js' 9 import { getIsRemoteMode } from '../../bootstrap/state.js' 10 import { getSystemPrompt } from '../../constants/prompts.js' 11 import { getSystemContext, getUserContext } from '../../context.js' 12 import type { CanUseToolFn } from '../../hooks/useCanUseTool.js' 13 import type { Tool, ToolUseContext } from '../../Tool.js' 14 import { FILE_EDIT_TOOL_NAME } from '../../tools/FileEditTool/constants.js' 15 import { 16 FileReadTool, 17 type Output as FileReadToolOutput, 18 } from '../../tools/FileReadTool/FileReadTool.js' 19 import type { Message } from '../../types/message.js' 20 import { count } from '../../utils/array.js' 21 import { 22 createCacheSafeParams, 23 createSubagentContext, 24 runForkedAgent, 25 } from '../../utils/forkedAgent.js' 26 import { getFsImplementation } from '../../utils/fsOperations.js' 27 import { 28 type REPLHookContext, 29 registerPostSamplingHook, 30 } from '../../utils/hooks/postSamplingHooks.js' 31 import { 32 createUserMessage, 33 hasToolCallsInLastAssistantTurn, 34 } from '../../utils/messages.js' 35 import { 36 getSessionMemoryDir, 37 getSessionMemoryPath, 38 } from '../../utils/permissions/filesystem.js' 39 import { sequential } from '../../utils/sequential.js' 40 import { asSystemPrompt } from '../../utils/systemPromptType.js' 41 import { getTokenUsage, tokenCountWithEstimation } from '../../utils/tokens.js' 42 import { logEvent } from '../analytics/index.js' 43 import { isAutoCompactEnabled } from '../compact/autoCompact.js' 44 import { 45 buildSessionMemoryUpdatePrompt, 46 loadSessionMemoryTemplate, 47 } from './prompts.js' 48 import { 49 DEFAULT_SESSION_MEMORY_CONFIG, 50 getSessionMemoryConfig, 51 getToolCallsBetweenUpdates, 52 hasMetInitializationThreshold, 53 hasMetUpdateThreshold, 54 isSessionMemoryInitialized, 55 markExtractionCompleted, 56 markExtractionStarted, 57 markSessionMemoryInitialized, 58 recordExtractionTokenCount, 59 type SessionMemoryConfig, 60 setLastSummarizedMessageId, 61 setSessionMemoryConfig, 62 } from './sessionMemoryUtils.js' 63 64 // ============================================================================ 65 // Feature Gate and Config (Cached - Non-blocking) 66 // ============================================================================ 67 // These functions return cached values from disk immediately without blocking 68 // on GrowthBook initialization. Values may be stale but are updated in background. 69 70 import { errorMessage, getErrnoCode } from '../../utils/errors.js' 71 import { 72 getDynamicConfig_CACHED_MAY_BE_STALE, 73 getFeatureValue_CACHED_MAY_BE_STALE, 74 } from '../analytics/growthbook.js' 75 76 /** 77 * Check if session memory feature is enabled. 78 * Uses cached gate value - returns immediately without blocking. 79 */ 80 function isSessionMemoryGateEnabled(): boolean { 81 return getFeatureValue_CACHED_MAY_BE_STALE('tengu_session_memory', false) 82 } 83 84 /** 85 * Get session memory config from cache. 86 * Returns immediately without blocking - value may be stale. 87 */ 88 function getSessionMemoryRemoteConfig(): Partial<SessionMemoryConfig> { 89 return getDynamicConfig_CACHED_MAY_BE_STALE<Partial<SessionMemoryConfig>>( 90 'tengu_sm_config', 91 {}, 92 ) 93 } 94 95 // ============================================================================ 96 // Module State 97 // ============================================================================ 98 99 let lastMemoryMessageUuid: string | undefined 100 101 /** 102 * Reset the last memory message UUID (for testing) 103 */ 104 export function resetLastMemoryMessageUuid(): void { 105 lastMemoryMessageUuid = undefined 106 } 107 108 function countToolCallsSince( 109 messages: Message[], 110 sinceUuid: string | undefined, 111 ): number { 112 let toolCallCount = 0 113 let foundStart = sinceUuid === null || sinceUuid === undefined 114 115 for (const message of messages) { 116 if (!foundStart) { 117 if (message.uuid === sinceUuid) { 118 foundStart = true 119 } 120 continue 121 } 122 123 if (message.type === 'assistant') { 124 const content = message.message.content 125 if (Array.isArray(content)) { 126 toolCallCount += count(content, block => block.type === 'tool_use') 127 } 128 } 129 } 130 131 return toolCallCount 132 } 133 134 export function shouldExtractMemory(messages: Message[]): boolean { 135 // Check if we've met the initialization threshold 136 // Uses total context window tokens (same as autocompact) for consistent behavior 137 const currentTokenCount = tokenCountWithEstimation(messages) 138 if (!isSessionMemoryInitialized()) { 139 if (!hasMetInitializationThreshold(currentTokenCount)) { 140 return false 141 } 142 markSessionMemoryInitialized() 143 } 144 145 // Check if we've met the minimum tokens between updates threshold 146 // Uses context window growth since last extraction (same metric as init threshold) 147 const hasMetTokenThreshold = hasMetUpdateThreshold(currentTokenCount) 148 149 // Check if we've met the tool calls threshold 150 const toolCallsSinceLastUpdate = countToolCallsSince( 151 messages, 152 lastMemoryMessageUuid, 153 ) 154 const hasMetToolCallThreshold = 155 toolCallsSinceLastUpdate >= getToolCallsBetweenUpdates() 156 157 // Check if the last assistant turn has no tool calls (safe to extract) 158 const hasToolCallsInLastTurn = hasToolCallsInLastAssistantTurn(messages) 159 160 // Trigger extraction when: 161 // 1. Both thresholds are met (tokens AND tool calls), OR 162 // 2. No tool calls in last turn AND token threshold is met 163 // (to ensure we extract at natural conversation breaks) 164 // 165 // IMPORTANT: The token threshold (minimumTokensBetweenUpdate) is ALWAYS required. 166 // Even if the tool call threshold is met, extraction won't happen until the 167 // token threshold is also satisfied. This prevents excessive extractions. 168 const shouldExtract = 169 (hasMetTokenThreshold && hasMetToolCallThreshold) || 170 (hasMetTokenThreshold && !hasToolCallsInLastTurn) 171 172 if (shouldExtract) { 173 const lastMessage = messages[messages.length - 1] 174 if (lastMessage?.uuid) { 175 lastMemoryMessageUuid = lastMessage.uuid 176 } 177 return true 178 } 179 180 return false 181 } 182 183 async function setupSessionMemoryFile( 184 toolUseContext: ToolUseContext, 185 ): Promise<{ memoryPath: string; currentMemory: string }> { 186 const fs = getFsImplementation() 187 188 // Set up directory and file 189 const sessionMemoryDir = getSessionMemoryDir() 190 await fs.mkdir(sessionMemoryDir, { mode: 0o700 }) 191 192 const memoryPath = getSessionMemoryPath() 193 194 // Create the memory file if it doesn't exist (wx = O_CREAT|O_EXCL) 195 try { 196 await writeFile(memoryPath, '', { 197 encoding: 'utf-8', 198 mode: 0o600, 199 flag: 'wx', 200 }) 201 // Only load template if file was just created 202 const template = await loadSessionMemoryTemplate() 203 await writeFile(memoryPath, template, { 204 encoding: 'utf-8', 205 mode: 0o600, 206 }) 207 } catch (e: unknown) { 208 const code = getErrnoCode(e) 209 if (code !== 'EEXIST') { 210 throw e 211 } 212 } 213 214 // Drop any cached entry so FileReadTool's dedup doesn't return a 215 // file_unchanged stub — we need the actual content. The Read repopulates it. 216 toolUseContext.readFileState.delete(memoryPath) 217 const result = await FileReadTool.call( 218 { file_path: memoryPath }, 219 toolUseContext, 220 ) 221 let currentMemory = '' 222 223 const output = result.data as FileReadToolOutput 224 if (output.type === 'text') { 225 currentMemory = output.file.content 226 } 227 228 logEvent('tengu_session_memory_file_read', { 229 content_length: currentMemory.length, 230 }) 231 232 return { memoryPath, currentMemory } 233 } 234 235 /** 236 * Initialize session memory config from remote config (lazy initialization). 237 * Memoized - only runs once per session, subsequent calls return immediately. 238 * Uses cached config values - non-blocking. 239 */ 240 const initSessionMemoryConfigIfNeeded = memoize((): void => { 241 // Load config from cache (non-blocking, may be stale) 242 const remoteConfig = getSessionMemoryRemoteConfig() 243 244 // Only use remote values if they are explicitly set (non-zero positive numbers) 245 // This ensures sensible defaults aren't overridden by zero values 246 const config: SessionMemoryConfig = { 247 minimumMessageTokensToInit: 248 remoteConfig.minimumMessageTokensToInit && 249 remoteConfig.minimumMessageTokensToInit > 0 250 ? remoteConfig.minimumMessageTokensToInit 251 : DEFAULT_SESSION_MEMORY_CONFIG.minimumMessageTokensToInit, 252 minimumTokensBetweenUpdate: 253 remoteConfig.minimumTokensBetweenUpdate && 254 remoteConfig.minimumTokensBetweenUpdate > 0 255 ? remoteConfig.minimumTokensBetweenUpdate 256 : DEFAULT_SESSION_MEMORY_CONFIG.minimumTokensBetweenUpdate, 257 toolCallsBetweenUpdates: 258 remoteConfig.toolCallsBetweenUpdates && 259 remoteConfig.toolCallsBetweenUpdates > 0 260 ? remoteConfig.toolCallsBetweenUpdates 261 : DEFAULT_SESSION_MEMORY_CONFIG.toolCallsBetweenUpdates, 262 } 263 setSessionMemoryConfig(config) 264 }) 265 266 /** 267 * Session memory post-sampling hook that extracts and updates session notes 268 */ 269 // Track if we've logged the gate check failure this session (to avoid spam) 270 let hasLoggedGateFailure = false 271 272 const extractSessionMemory = sequential(async function ( 273 context: REPLHookContext, 274 ): Promise<void> { 275 const { messages, toolUseContext, querySource } = context 276 277 // Only run session memory on main REPL thread 278 if (querySource !== 'repl_main_thread') { 279 // Don't log this - it's expected for subagents, teammates, etc. 280 return 281 } 282 283 // Check gate lazily when hook runs (cached, non-blocking) 284 if (!isSessionMemoryGateEnabled()) { 285 // Log gate failure once per session (ant-only) 286 if (process.env.USER_TYPE === 'ant' && !hasLoggedGateFailure) { 287 hasLoggedGateFailure = true 288 logEvent('tengu_session_memory_gate_disabled', {}) 289 } 290 return 291 } 292 293 // Initialize config from remote (lazy, only once) 294 initSessionMemoryConfigIfNeeded() 295 296 if (!shouldExtractMemory(messages)) { 297 return 298 } 299 300 markExtractionStarted() 301 302 // Create isolated context for setup to avoid polluting parent's cache 303 const setupContext = createSubagentContext(toolUseContext) 304 305 // Set up file system and read current state with isolated context 306 const { memoryPath, currentMemory } = 307 await setupSessionMemoryFile(setupContext) 308 309 // Create extraction message 310 const userPrompt = await buildSessionMemoryUpdatePrompt( 311 currentMemory, 312 memoryPath, 313 ) 314 315 // Run session memory extraction using runForkedAgent for prompt caching 316 // runForkedAgent creates an isolated context to prevent mutation of parent state 317 // Pass setupContext.readFileState so the forked agent can edit the memory file 318 await runForkedAgent({ 319 promptMessages: [createUserMessage({ content: userPrompt })], 320 cacheSafeParams: createCacheSafeParams(context), 321 canUseTool: createMemoryFileCanUseTool(memoryPath), 322 querySource: 'session_memory', 323 forkLabel: 'session_memory', 324 overrides: { readFileState: setupContext.readFileState }, 325 }) 326 327 // Log extraction event for tracking frequency 328 // Use the token usage from the last message in the conversation 329 const lastMessage = messages[messages.length - 1] 330 const usage = lastMessage ? getTokenUsage(lastMessage) : undefined 331 const config = getSessionMemoryConfig() 332 logEvent('tengu_session_memory_extraction', { 333 input_tokens: usage?.input_tokens, 334 output_tokens: usage?.output_tokens, 335 cache_read_input_tokens: usage?.cache_read_input_tokens ?? undefined, 336 cache_creation_input_tokens: 337 usage?.cache_creation_input_tokens ?? undefined, 338 config_min_message_tokens_to_init: config.minimumMessageTokensToInit, 339 config_min_tokens_between_update: config.minimumTokensBetweenUpdate, 340 config_tool_calls_between_updates: config.toolCallsBetweenUpdates, 341 }) 342 343 // Record the context size at extraction for tracking minimumTokensBetweenUpdate 344 recordExtractionTokenCount(tokenCountWithEstimation(messages)) 345 346 // Update lastSummarizedMessageId after successful completion 347 updateLastSummarizedMessageIdIfSafe(messages) 348 349 markExtractionCompleted() 350 }) 351 352 /** 353 * Initialize session memory by registering the post-sampling hook. 354 * This is synchronous to avoid race conditions during startup. 355 * The gate check and config loading happen lazily when the hook runs. 356 */ 357 export function initSessionMemory(): void { 358 if (getIsRemoteMode()) return 359 // Session memory is used for compaction, so respect auto-compact settings 360 const autoCompactEnabled = isAutoCompactEnabled() 361 362 // Log initialization state (ant-only to avoid noise in external logs) 363 if (process.env.USER_TYPE === 'ant') { 364 logEvent('tengu_session_memory_init', { 365 auto_compact_enabled: autoCompactEnabled, 366 }) 367 } 368 369 if (!autoCompactEnabled) { 370 return 371 } 372 373 // Register hook unconditionally - gate check happens lazily when hook runs 374 registerPostSamplingHook(extractSessionMemory) 375 } 376 377 export type ManualExtractionResult = { 378 success: boolean 379 memoryPath?: string 380 error?: string 381 } 382 383 /** 384 * Manually trigger session memory extraction, bypassing threshold checks. 385 * Used by the /summary command. 386 */ 387 export async function manuallyExtractSessionMemory( 388 messages: Message[], 389 toolUseContext: ToolUseContext, 390 ): Promise<ManualExtractionResult> { 391 if (messages.length === 0) { 392 return { success: false, error: 'No messages to summarize' } 393 } 394 markExtractionStarted() 395 396 try { 397 // Create isolated context for setup to avoid polluting parent's cache 398 const setupContext = createSubagentContext(toolUseContext) 399 400 // Set up file system and read current state with isolated context 401 const { memoryPath, currentMemory } = 402 await setupSessionMemoryFile(setupContext) 403 404 // Create extraction message 405 const userPrompt = await buildSessionMemoryUpdatePrompt( 406 currentMemory, 407 memoryPath, 408 ) 409 410 // Get system prompt for cache-safe params 411 const { tools, mainLoopModel } = toolUseContext.options 412 const [rawSystemPrompt, userContext, systemContext] = await Promise.all([ 413 getSystemPrompt(tools, mainLoopModel), 414 getUserContext(), 415 getSystemContext(), 416 ]) 417 const systemPrompt = asSystemPrompt(rawSystemPrompt) 418 419 // Run session memory extraction using runForkedAgent 420 await runForkedAgent({ 421 promptMessages: [createUserMessage({ content: userPrompt })], 422 cacheSafeParams: { 423 systemPrompt, 424 userContext, 425 systemContext, 426 toolUseContext: setupContext, 427 forkContextMessages: messages, 428 }, 429 canUseTool: createMemoryFileCanUseTool(memoryPath), 430 querySource: 'session_memory', 431 forkLabel: 'session_memory_manual', 432 overrides: { readFileState: setupContext.readFileState }, 433 }) 434 435 // Log manual extraction event 436 logEvent('tengu_session_memory_manual_extraction', {}) 437 438 // Record the context size at extraction for tracking minimumTokensBetweenUpdate 439 recordExtractionTokenCount(tokenCountWithEstimation(messages)) 440 441 // Update lastSummarizedMessageId after successful completion 442 updateLastSummarizedMessageIdIfSafe(messages) 443 444 return { success: true, memoryPath } 445 } catch (error) { 446 return { 447 success: false, 448 error: errorMessage(error), 449 } 450 } finally { 451 markExtractionCompleted() 452 } 453 } 454 455 // Helper functions 456 457 /** 458 * Creates a canUseTool function that only allows Edit for the exact memory file. 459 */ 460 export function createMemoryFileCanUseTool(memoryPath: string): CanUseToolFn { 461 return async (tool: Tool, input: unknown) => { 462 if ( 463 tool.name === FILE_EDIT_TOOL_NAME && 464 typeof input === 'object' && 465 input !== null && 466 'file_path' in input 467 ) { 468 const filePath = input.file_path 469 if (typeof filePath === 'string' && filePath === memoryPath) { 470 return { behavior: 'allow' as const, updatedInput: input } 471 } 472 } 473 return { 474 behavior: 'deny' as const, 475 message: `only ${FILE_EDIT_TOOL_NAME} on ${memoryPath} is allowed`, 476 decisionReason: { 477 type: 'other' as const, 478 reason: `only ${FILE_EDIT_TOOL_NAME} on ${memoryPath} is allowed`, 479 }, 480 } 481 } 482 } 483 484 /** 485 * Updates lastSummarizedMessageId after successful extraction. 486 * Only sets it if the last message doesn't have tool calls (to avoid orphaned tool_results). 487 */ 488 function updateLastSummarizedMessageIdIfSafe(messages: Message[]): void { 489 if (!hasToolCallsInLastAssistantTurn(messages)) { 490 const lastMessage = messages[messages.length - 1] 491 if (lastMessage?.uuid) { 492 setLastSummarizedMessageId(lastMessage.uuid) 493 } 494 } 495 }