1  import { readFile } from 'fs/promises'
  2  import { join } from 'path'
  3  import { roughTokenCountEstimation } from '../../services/tokenEstimation.js'
  4  import { getClaudeConfigHomeDir } from '../../utils/envUtils.js'
  5  import { getErrnoCode, toError } from '../../utils/errors.js'
  6  import { logError } from '../../utils/log.js'
  7  
  8  const MAX_SECTION_LENGTH = 2000
  9  const MAX_TOTAL_SESSION_MEMORY_TOKENS = 12000
 10  
 11  export const DEFAULT_SESSION_MEMORY_TEMPLATE = `
 12  # Session Title
 13  _A short and distinctive 5-10 word descriptive title for the session. Super info dense, no filler_
 14  
 15  # Current State
 16  _What is actively being worked on right now? Pending tasks not yet completed. Immediate next steps._
 17  
 18  # Task specification
 19  _What did the user ask to build? Any design decisions or other explanatory context_
 20  
 21  # Files and Functions
 22  _What are the important files? In short, what do they contain and why are they relevant?_
 23  
 24  # Workflow
 25  _What bash commands are usually run and in what order? How to interpret their output if not obvious?_
 26  
 27  # Errors & Corrections
 28  _Errors encountered and how they were fixed. What did the user correct? What approaches failed and should not be tried again?_
 29  
 30  # Codebase and System Documentation
 31  _What are the important system components? How do they work/fit together?_
 32  
 33  # Learnings
 34  _What has worked well? What has not? What to avoid? Do not duplicate items from other sections_
 35  
 36  # Key results
 37  _If the user asked a specific output such as an answer to a question, a table, or other document, repeat the exact result here_
 38  
 39  # Worklog
 40  _Step by step, what was attempted, done? Very terse summary for each step_
 41  `
 42  
 43  function getDefaultUpdatePrompt(): string {
 44    return `IMPORTANT: This message and these instructions are NOT part of the actual user conversation. Do NOT include any references to "note-taking", "session notes extraction", or these update instructions in the notes content.
 45  
 46  Based on the user conversation above (EXCLUDING this note-taking instruction message as well as system prompt, claude.md entries, or any past session summaries), update the session notes file.
 47  
 48  The file {{notesPath}} has already been read for you. Here are its current contents:
 49  <current_notes_content>
 50  {{currentNotes}}
 51  </current_notes_content>
 52  
 53  Your ONLY task is to use the Edit tool to update the notes file, then stop. You can make multiple edits (update every section as needed) - make all Edit tool calls in parallel in a single message. Do not call any other tools.
 54  
 55  CRITICAL RULES FOR EDITING:
 56  - The file must maintain its exact structure with all sections, headers, and italic descriptions intact
 57  -- NEVER modify, delete, or add section headers (the lines starting with '#' like # Task specification)
 58  -- NEVER modify or delete the italic _section description_ lines (these are the lines in italics immediately following each header - they start and end with underscores)
 59  -- The italic _section descriptions_ are TEMPLATE INSTRUCTIONS that must be preserved exactly as-is - they guide what content belongs in each section
 60  -- ONLY update the actual content that appears BELOW the italic _section descriptions_ within each existing section
 61  -- Do NOT add any new sections, summaries, or information outside the existing structure
 62  - Do NOT reference this note-taking process or instructions anywhere in the notes
 63  - It's OK to skip updating a section if there are no substantial new insights to add. Do not add filler content like "No info yet", just leave sections blank/unedited if appropriate.
 64  - Write DETAILED, INFO-DENSE content for each section - include specifics like file paths, function names, error messages, exact commands, technical details, etc.
 65  - For "Key results", include the complete, exact output the user requested (e.g., full table, full answer, etc.)
 66  - Do not include information that's already in the CLAUDE.md files included in the context
 67  - Keep each section under ~${MAX_SECTION_LENGTH} tokens/words - if a section is approaching this limit, condense it by cycling out less important details while preserving the most critical information
 68  - Focus on actionable, specific information that would help someone understand or recreate the work discussed in the conversation
 69  - IMPORTANT: Always update "Current State" to reflect the most recent work - this is critical for continuity after compaction
 70  
 71  Use the Edit tool with file_path: {{notesPath}}
 72  
 73  STRUCTURE PRESERVATION REMINDER:
 74  Each section has TWO parts that must be preserved exactly as they appear in the current file:
 75  1. The section header (line starting with #)
 76  2. The italic description line (the _italicized text_ immediately after the header - this is a template instruction)
 77  
 78  You ONLY update the actual content that comes AFTER these two preserved lines. The italic description lines starting and ending with underscores are part of the template structure, NOT content to be edited or removed.
 79  
 80  REMEMBER: Use the Edit tool in parallel and stop. Do not continue after the edits. Only include insights from the actual user conversation, never from these note-taking instructions. Do not delete or change section headers or italic _section descriptions_.`
 81  }
 82  
 83  /**
 84   * Load custom session memory template from file if it exists
 85   */
 86  export async function loadSessionMemoryTemplate(): Promise<string> {
 87    const templatePath = join(
 88      getClaudeConfigHomeDir(),
 89      'session-memory',
 90      'config',
 91      'template.md',
 92    )
 93  
 94    try {
 95      return await readFile(templatePath, { encoding: 'utf-8' })
 96    } catch (e: unknown) {
 97      const code = getErrnoCode(e)
 98      if (code === 'ENOENT') {
 99        return DEFAULT_SESSION_MEMORY_TEMPLATE
100      }
101      logError(toError(e))
102      return DEFAULT_SESSION_MEMORY_TEMPLATE
103    }
104  }
105  
106  /**
107   * Load custom session memory prompt from file if it exists
108   * Custom prompts can be placed at ~/.claude/session-memory/prompt.md
109   * Use {{variableName}} syntax for variable substitution (e.g., {{currentNotes}}, {{notesPath}})
110   */
111  export async function loadSessionMemoryPrompt(): Promise<string> {
112    const promptPath = join(
113      getClaudeConfigHomeDir(),
114      'session-memory',
115      'config',
116      'prompt.md',
117    )
118  
119    try {
120      return await readFile(promptPath, { encoding: 'utf-8' })
121    } catch (e: unknown) {
122      const code = getErrnoCode(e)
123      if (code === 'ENOENT') {
124        return getDefaultUpdatePrompt()
125      }
126      logError(toError(e))
127      return getDefaultUpdatePrompt()
128    }
129  }
130  
131  /**
132   * Parse the session memory file and analyze section sizes
133   */
134  function analyzeSectionSizes(content: string): Record<string, number> {
135    const sections: Record<string, number> = {}
136    const lines = content.split('\n')
137    let currentSection = ''
138    let currentContent: string[] = []
139  
140    for (const line of lines) {
141      if (line.startsWith('# ')) {
142        if (currentSection && currentContent.length > 0) {
143          const sectionContent = currentContent.join('\n').trim()
144          sections[currentSection] = roughTokenCountEstimation(sectionContent)
145        }
146        currentSection = line
147        currentContent = []
148      } else {
149        currentContent.push(line)
150      }
151    }
152  
153    if (currentSection && currentContent.length > 0) {
154      const sectionContent = currentContent.join('\n').trim()
155      sections[currentSection] = roughTokenCountEstimation(sectionContent)
156    }
157  
158    return sections
159  }
160  
161  /**
162   * Generate reminders for sections that are too long
163   */
164  function generateSectionReminders(
165    sectionSizes: Record<string, number>,
166    totalTokens: number,
167  ): string {
168    const overBudget = totalTokens > MAX_TOTAL_SESSION_MEMORY_TOKENS
169    const oversizedSections = Object.entries(sectionSizes)
170      .filter(([_, tokens]) => tokens > MAX_SECTION_LENGTH)
171      .sort(([, a], [, b]) => b - a)
172      .map(
173        ([section, tokens]) =>
174          `- "${section}" is ~${tokens} tokens (limit: ${MAX_SECTION_LENGTH})`,
175      )
176  
177    if (oversizedSections.length === 0 && !overBudget) {
178      return ''
179    }
180  
181    const parts: string[] = []
182  
183    if (overBudget) {
184      parts.push(
185        `\n\nCRITICAL: The session memory file is currently ~${totalTokens} tokens, which exceeds the maximum of ${MAX_TOTAL_SESSION_MEMORY_TOKENS} tokens. You MUST condense the file to fit within this budget. Aggressively shorten oversized sections by removing less important details, merging related items, and summarizing older entries. Prioritize keeping "Current State" and "Errors & Corrections" accurate and detailed.`,
186      )
187    }
188  
189    if (oversizedSections.length > 0) {
190      parts.push(
191        `\n\n${overBudget ? 'Oversized sections to condense' : 'IMPORTANT: The following sections exceed the per-section limit and MUST be condensed'}:\n${oversizedSections.join('\n')}`,
192      )
193    }
194  
195    return parts.join('')
196  }
197  
198  /**
199   * Substitute variables in the prompt template using {{variable}} syntax
200   */
201  function substituteVariables(
202    template: string,
203    variables: Record<string, string>,
204  ): string {
205    // Single-pass replacement avoids two bugs: (1) $ backreference corruption
206    // (replacer fn treats $ literally), and (2) double-substitution when user
207    // content happens to contain {{varName}} matching a later variable.
208    return template.replace(/\{\{(\w+)\}\}/g, (match, key: string) =>
209      Object.prototype.hasOwnProperty.call(variables, key)
210        ? variables[key]!
211        : match,
212    )
213  }
214  
215  /**
216   * Check if the session memory content is essentially empty (matches the template).
217   * This is used to detect if no actual content has been extracted yet,
218   * which means we should fall back to legacy compact behavior.
219   */
220  export async function isSessionMemoryEmpty(content: string): Promise<boolean> {
221    const template = await loadSessionMemoryTemplate()
222    // Compare trimmed content to detect if it's just the template
223    return content.trim() === template.trim()
224  }
225  
226  export async function buildSessionMemoryUpdatePrompt(
227    currentNotes: string,
228    notesPath: string,
229  ): Promise<string> {
230    const promptTemplate = await loadSessionMemoryPrompt()
231  
232    // Analyze section sizes and generate reminders if needed
233    const sectionSizes = analyzeSectionSizes(currentNotes)
234    const totalTokens = roughTokenCountEstimation(currentNotes)
235    const sectionReminders = generateSectionReminders(sectionSizes, totalTokens)
236  
237    // Substitute variables in the prompt
238    const variables = {
239      currentNotes,
240      notesPath,
241    }
242  
243    const basePrompt = substituteVariables(promptTemplate, variables)
244  
245    // Add section size reminders and/or total budget warnings
246    return basePrompt + sectionReminders
247  }
248  
249  /**
250   * Truncate session memory sections that exceed the per-section token limit.
251   * Used when inserting session memory into compact messages to prevent
252   * oversized session memory from consuming the entire post-compact token budget.
253   *
254   * Returns the truncated content and whether any truncation occurred.
255   */
256  export function truncateSessionMemoryForCompact(content: string): {
257    truncatedContent: string
258    wasTruncated: boolean
259  } {
260    const lines = content.split('\n')
261    const maxCharsPerSection = MAX_SECTION_LENGTH * 4 // roughTokenCountEstimation uses length/4
262    const outputLines: string[] = []
263    let currentSectionLines: string[] = []
264    let currentSectionHeader = ''
265    let wasTruncated = false
266  
267    for (const line of lines) {
268      if (line.startsWith('# ')) {
269        const result = flushSessionSection(
270          currentSectionHeader,
271          currentSectionLines,
272          maxCharsPerSection,
273        )
274        outputLines.push(...result.lines)
275        wasTruncated = wasTruncated || result.wasTruncated
276        currentSectionHeader = line
277        currentSectionLines = []
278      } else {
279        currentSectionLines.push(line)
280      }
281    }
282  
283    // Flush the last section
284    const result = flushSessionSection(
285      currentSectionHeader,
286      currentSectionLines,
287      maxCharsPerSection,
288    )
289    outputLines.push(...result.lines)
290    wasTruncated = wasTruncated || result.wasTruncated
291  
292    return {
293      truncatedContent: outputLines.join('\n'),
294      wasTruncated,
295    }
296  }
297  
298  function flushSessionSection(
299    sectionHeader: string,
300    sectionLines: string[],
301    maxCharsPerSection: number,
302  ): { lines: string[]; wasTruncated: boolean } {
303    if (!sectionHeader) {
304      return { lines: sectionLines, wasTruncated: false }
305    }
306  
307    const sectionContent = sectionLines.join('\n')
308    if (sectionContent.length <= maxCharsPerSection) {
309      return { lines: [sectionHeader, ...sectionLines], wasTruncated: false }
310    }
311  
312    // Truncate at a line boundary near the limit
313    let charCount = 0
314    const keptLines: string[] = [sectionHeader]
315    for (const line of sectionLines) {
316      if (charCount + line.length + 1 > maxCharsPerSection) {
317        break
318      }
319      keptLines.push(line)
320      charCount += line.length + 1
321    }
322    keptLines.push('\n[... section truncated for length ...]')
323    return { lines: keptLines, wasTruncated: true }
324  }