| 1 | import fs from 'fs' |
| 2 | import path from 'path' |
| 3 | import log from 'electron-log/main.js' |
| 4 | import type { ThinkingChatMessage, ThinkingStage } from '@shared/thinking' |
| 5 | import { getStagePrompt } from './prompts' |
| 6 | import { VALID_TRANSITIONS } from './stage-manager' |
| 7 | |
| 8 | const MAX_INLINE_THINKING_CHARS = 12_000 |
| 9 | const MAX_INLINE_THINKING_PAGES = 24 |
| 10 | const MAX_THINKING_MAP_PAGES = 160 |
| 11 | |
| 12 | export interface ThinkingContextArgs { |
| 13 | stage: ThinkingStage |
| 14 | thinkingMd: string |
| 15 | contextMd: string |
| 16 | sourcesDir: string |
| 17 | userMessage: string |
| 18 | recentMessages?: ThinkingChatMessage[] |
| 19 | } |
| 20 | |
| 21 | function countThinkingPages(markdown: string): number { |
| 22 | const matches = markdown.match(/^##\s*Page\s+\d+\s*:/gm) |
| 23 | return matches ? matches.length : 0 |
| 24 | } |
| 25 | |
| 26 | function readSection(markdown: string, heading: string): string { |
| 27 | const escaped = heading.replace(/[.*+?^${}()|[\]\\]/g, '\\$&') |
| 28 | const inline = markdown.match(new RegExp(`^##\\s*${escaped}\\s*:\\s*(.+)`, 'm')) |
| 29 | if (inline?.[1]?.trim()) return inline[1].trim() |
| 30 | const block = markdown.match( |
| 31 | new RegExp(`^##\\s*${escaped}\\s*\\n([\\s\\S]*?)(?=^##\\s+|$)`, 'm') |
| 32 | ) |
| 33 | return block?.[1]?.trim() || '' |
| 34 | } |
| 35 | |
| 36 | function buildThinkingPageMap(markdown: string): string { |
| 37 | const pageHeadings = Array.from(markdown.matchAll(/^##\s*Page\s+(\d+)\s*:\s*(.+)$/gm)) |
| 38 | const lines = pageHeadings.slice(0, MAX_THINKING_MAP_PAGES).map((match) => { |
| 39 | const pageNumber = Number.parseInt(match[1], 10) |
| 40 | const title = match[2].trim() |
| 41 | return Number.isFinite(pageNumber) && title ? `- Page ${pageNumber}: ${title}` : '' |
| 42 | }).filter(Boolean) |
| 43 | |
| 44 | return [ |
| 45 | '# Thinking Brief Summary', |
| 46 | readSection(markdown, 'Topic') ? `## Topic\n${readSection(markdown, 'Topic')}` : '', |
| 47 | readSection(markdown, 'Audience') ? `## Audience\n${readSection(markdown, 'Audience')}` : '', |
| 48 | readSection(markdown, 'Setting') ? `## Setting\n${readSection(markdown, 'Setting')}` : '', |
| 49 | readSection(markdown, 'Tone') ? `## Tone\n${readSection(markdown, 'Tone')}` : '', |
| 50 | readSection(markdown, 'Style') ? `## Style\n${readSection(markdown, 'Style')}` : '', |
| 51 | readSection(markdown, 'Font') ? `## Font\n${readSection(markdown, 'Font')}` : '', |
| 52 | readSection(markdown, 'Page Count') ? `## Page Count\n${readSection(markdown, 'Page Count')}` : '', |
| 53 | lines.length > 0 ? `## Page Map\n${lines.join('\n')}` : '', |
| 54 | pageHeadings.length > lines.length |
| 55 | ? `\nOnly the first ${lines.length} page headings are shown. Use update_thinking_document with pageStart to modify page ranges instead of reading or rewriting the full thinking.md.` |
| 56 | : '\nUse update_thinking_document with pageStart to modify page ranges instead of reading or rewriting the full thinking.md.' |
| 57 | ] |
| 58 | .filter(Boolean) |
| 59 | .join('\n\n') |
| 60 | } |
| 61 | |
| 62 | function buildCurrentThinkingContext(thinkingMd: string): string { |
| 63 | const trimmed = thinkingMd.trim() |
| 64 | if (!trimmed) return '' |
| 65 | const pageCount = countThinkingPages(trimmed) |
| 66 | if (trimmed.length <= MAX_INLINE_THINKING_CHARS && pageCount <= MAX_INLINE_THINKING_PAGES) { |
| 67 | return trimmed |
| 68 | } |
| 69 | return buildThinkingPageMap(trimmed) |
| 70 | } |
| 71 | |
| 72 | export async function buildThinkingContext(args: ThinkingContextArgs): Promise<{ |
| 73 | systemPrompt: string |
| 74 | userMessage: string |
| 75 | sourceContent: string |
| 76 | }> { |
| 77 | const { stage, thinkingMd, contextMd, sourcesDir, userMessage, recentMessages } = args |
| 78 | |
| 79 | const stagePrompt = getStagePrompt(stage) |
| 80 | const validTargets = VALID_TRANSITIONS[stage].filter((s) => s !== stage) |
| 81 | const stageAwareSuffix = |
| 82 | validTargets.length > 0 |
| 83 | ? `\n\nYou are in stage "${stage}". When the user's intent clearly matches a later stage, call update_context_document with \`stage\` set to the target stage. Valid transitions from ${stage}: ${validTargets.join(', ')}.` |
| 84 | : '' |
| 85 | const systemPrompt = stagePrompt + stageAwareSuffix |
| 86 | |
| 87 | // Build source file index instead of inlining content — AI will use read_file/grep tools to read on demand |
| 88 | let sourceContent = '' |
| 89 | if (fs.existsSync(sourcesDir)) { |
| 90 | const entries = await fs.promises.readdir(sourcesDir) |
| 91 | const fileEntries: string[] = [] |
| 92 | for (const entry of entries) { |
| 93 | const filePath = path.join(sourcesDir, entry) |
| 94 | try { |
| 95 | const stat = await fs.promises.stat(filePath) |
| 96 | if (!stat.isFile()) continue |
| 97 | fileEntries.push(`- /sources/${entry}`) |
| 98 | } catch { |
| 99 | // skip unreadable files |
| 100 | } |
| 101 | } |
| 102 | if (fileEntries.length > 0) { |
| 103 | sourceContent = fileEntries.join('\n') |
| 104 | } |
| 105 | } |
| 106 | |
| 107 | const contextParts: string[] = [] |
| 108 | |
| 109 | const currentThinkingContext = buildCurrentThinkingContext(thinkingMd) |
| 110 | if (currentThinkingContext) { |
| 111 | contextParts.push(`## Current Thinking Brief\n${currentThinkingContext}`) |
| 112 | } |
| 113 | |
| 114 | if (contextMd.trim()) { |
| 115 | contextParts.push(`## Context\n${contextMd}`) |
| 116 | } |
| 117 | |
| 118 | if (sourceContent) { |
| 119 | contextParts.push( |
| 120 | [ |
| 121 | '## Available Source Files', |
| 122 | 'The following source files are available.', |
| 123 | 'Use grep first, then read_file with a small offset/limit around relevant lines. For large files, build the outline incrementally from source sections instead of reading the whole file in one pass.', |
| 124 | sourceContent |
| 125 | ].join('\n') |
| 126 | ) |
| 127 | } else { |
| 128 | contextParts.push( |
| 129 | [ |
| 130 | '## Source Files', |
| 131 | 'No source files are available for this turn.', |
| 132 | 'Do not call read_file, grep/search, glob, or ls. Work only from the current thinking brief, context, recent conversation, and user message.' |
| 133 | ].join('\n') |
| 134 | ) |
| 135 | } |
| 136 | |
| 137 | const recentConversation = Array.isArray(recentMessages) |
| 138 | ? recentMessages |
| 139 | .slice(-8) |
| 140 | .map((message) => { |
| 141 | const role = message.role === 'assistant' ? 'Assistant' : 'User' |
| 142 | return `${role}: ${message.content.trim()}` |
| 143 | }) |
| 144 | .filter((line) => line.trim().length > 0) |
| 145 | .join('\n\n') |
| 146 | : '' |
| 147 | |
| 148 | if (recentConversation) { |
| 149 | contextParts.push(`## Recent Conversation\n${recentConversation}`) |
| 150 | } |
| 151 | |
| 152 | contextParts.push(`## User Message\n${userMessage}`) |
| 153 | |
| 154 | const fullUserMessage = contextParts.join('\n\n') |
| 155 | |
| 156 | log.info('[thinking:context] built', { |
| 157 | stage, |
| 158 | hasThinkingMd: thinkingMd.trim().length > 0, |
| 159 | thinkingMdLength: thinkingMd.trim().length, |
| 160 | thinkingPageCount: countThinkingPages(thinkingMd), |
| 161 | thinkingContextLength: currentThinkingContext.length, |
| 162 | hasSources: sourceContent.length > 0, |
| 163 | recentMessages: recentMessages?.length || 0, |
| 164 | messageLength: fullUserMessage.length |
| 165 | }) |
| 166 | |
| 167 | return { |
| 168 | systemPrompt, |
| 169 | userMessage: fullUserMessage, |
| 170 | sourceContent |
| 171 | } |
| 172 | } |
| 173 |