fix(chat): bound PDF context with retrieved chunks

apps/electron/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@proma/electron",
-  "version": "0.15.7",
+  "version": "0.15.8",
   "description": "Proma next gen ai software with general agents - Electron App",
   "main": "dist/main.cjs",
   "author": {
@@ -38,6 +38,7 @@
     "dist:debug": "bun run scripts/dist.ts --current-arch --verbose"
   },
   "dependencies": {
+    "@anthropic-ai/tokenizer": "0.0.4",
     "@anthropic-ai/claude-agent-sdk": "0.3.201",
     "@anthropic-ai/sdk": "^0.93.0",
     "@earendil-works/pi-agent-core": "0.80.9",

apps/electron/src/main/lib/agent-session-manager.test.ts

@@ -2,6 +2,7 @@ import { afterAll, beforeAll, describe, expect, mock, test } from 'bun:test'
 import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs'
 import * as os from 'node:os'
 import { join } from 'node:path'
+import { createElectronMock } from './test-helpers/electron-mock'
 
 type AgentSessionManager = typeof import('./agent-session-manager')
 
@@ -11,26 +12,9 @@ const originalHome = process.env.HOME
 const originalPromaDev = process.env.PROMA_DEV
 const originalClaudeConfigDir = process.env.CLAUDE_CONFIG_DIR
 
-mock.module('electron', () => ({
-  app: {
-    isPackaged: true,
-    getPath: () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'),
-  },
-  BrowserWindow: class {},
-  clipboard: {},
-  dialog: {},
-  nativeImage: { createFromPath: () => ({}) },
-  nativeTheme: {},
-  powerMonitor: {},
-  powerSaveBlocker: {},
-  screen: {},
-  shell: {},
-  safeStorage: {
-    isEncryptionAvailable: () => false,
-    encryptString: (value: string) => Buffer.from(value),
-    decryptString: (value: Buffer) => value.toString('utf-8'),
-  },
-}))
+mock.module('electron', () => createElectronMock(
+  () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'),
+))
 
 mock.module('node:os', () => ({
   ...os,

apps/electron/src/main/lib/agent-workspace-manager.test.ts

@@ -2,6 +2,7 @@ import { afterAll, beforeAll, beforeEach, describe, expect, mock, test } from 'b
 import { existsSync, mkdirSync, mkdtempSync, readdirSync, rmSync, symlinkSync, writeFileSync } from 'node:fs'
 import * as os from 'node:os'
 import { join } from 'node:path'
+import { createElectronMock } from './test-helpers/electron-mock'
 
 type AgentWorkspaceManager = typeof import('./agent-workspace-manager')
 type ConfigPathsModule = typeof import('./config-paths')
@@ -12,17 +13,9 @@ let tempHome: string
 const originalHome = process.env.HOME
 const originalPromaDev = process.env.PROMA_DEV
 
-mock.module('electron', () => ({
-  app: {
-    isPackaged: true,
-    getPath: () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'),
-  },
-  safeStorage: {
-    isEncryptionAvailable: () => false,
-    encryptString: (value: string) => Buffer.from(value),
-    decryptString: (value: Buffer) => value.toString('utf-8'),
-  },
-}))
+mock.module('electron', () => createElectronMock(
+  () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'),
+))
 
 mock.module('node:os', () => ({
   ...os,

apps/electron/src/main/lib/channel-runtime-api-key.test.ts

@@ -3,6 +3,7 @@ import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs'
 import * as os from 'node:os'
 import { join } from 'node:path'
 import { serializeCodexCredentials } from '@proma/shared'
+import { createElectronMock } from './test-helpers/electron-mock'
 
 type ChannelManagerModule = typeof import('./channel-manager')
 
@@ -11,20 +12,9 @@ let tempHome: string
 const originalHome = process.env.HOME
 const originalPromaDev = process.env.PROMA_DEV
 
-mock.module('electron', () => ({
-  app: {
-    isPackaged: true,
-    getPath: () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'),
-  },
-  safeStorage: {
-    isEncryptionAvailable: () => false,
-    encryptString: (value: string) => Buffer.from(value),
-    decryptString: (value: Buffer) => value.toString('utf-8'),
-  },
-  shell: {
-    openExternal: async () => undefined,
-  },
-}))
+mock.module('electron', () => createElectronMock(
+  () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'),
+))
 
 mock.module('node:os', () => ({
   ...os,

apps/electron/src/main/lib/chat-service.ts

@@ -21,11 +21,17 @@ import {
   streamSSE,
   fetchTitle,
 } from '@proma/core'
-import type { ImageAttachmentData, ContinuationMessage } from '@proma/core'
+import type { ImageAttachmentData, DocumentAttachmentData, ContinuationMessage } from '@proma/core'
 import { listChannels, resolveChannelRuntimeApiKey } from './channel-manager'
 import { appendMessage, updateConversationMeta, getConversationMessages } from './conversation-manager'
 import { readAttachmentAsBase64, isImageAttachment } from './attachment-service'
 import { extractTextFromAttachment, isDocumentAttachment } from './document-parser'
+import {
+  estimateDocumentTokensConservatively,
+  estimatePromptTokens,
+  PROMPT_SAFE_INPUT_TOKENS,
+  retrieveDocumentContexts,
+} from './document-context'
 import { getFetchFn } from './proxy-fetch'
 import { getEffectiveProxyUrl } from './proxy-settings-service'
 import { getEnabledTools } from './chat-tool-registry'
@@ -38,6 +44,18 @@ const activeControllers = new Map<string, AbortController>()
 /** 最大工具续接轮数（安全上限，防止极端情况下的无限循环） */
 const MAX_TOOL_ROUNDS = 999
 
+const NATIVE_DOCUMENT_MAX_FILE_BYTES = 8 * 1024 * 1024
+const NATIVE_DOCUMENT_MAX_TOTAL_BYTES = 16 * 1024 * 1024
+const NATIVE_DOCUMENT_MAX_TEXT_TOKENS = 80_000
+const NATIVE_DOCUMENT_MAX_TOTAL_TEXT_TOKENS = 120_000
+
+interface PreparedDocuments {
+  nativeDocumentIds: Set<string>
+  nativeEstimatedTokens: number
+  retrievalSources: Array<{ filename: string; text: string }>
+  notices: string[]
+}
+
 // ===== 平台相关：图片附件读取器 =====
 
 /**
@@ -57,6 +75,57 @@ function getImageAttachmentData(attachments?: FileAttachment[]): ImageAttachment
     }))
 }
 
+function getDocumentAttachmentData(
+  attachments: FileAttachment[] | undefined,
+  nativeDocumentIds: Set<string>,
+): DocumentAttachmentData[] {
+  return (attachments ?? []).filter((attachment) => nativeDocumentIds.has(attachment.id)).map((attachment) => ({
+    filename: attachment.filename,
+    mediaType: attachment.mediaType,
+    data: readAttachmentAsBase64(attachment.localPath),
+  }))
+}
+
+async function prepareDocuments(attachments?: FileAttachment[]): Promise<PreparedDocuments> {
+  const prepared: PreparedDocuments = {
+    nativeDocumentIds: new Set<string>(),
+    nativeEstimatedTokens: 0,
+    retrievalSources: [],
+    notices: [],
+  }
+  if (!attachments || attachments.length === 0) return prepared
+
+  let nativeBytes = 0
+  for (const attachment of attachments.filter((item) => isDocumentAttachment(item.mediaType))) {
+    try {
+      const text = await extractTextFromAttachment(attachment.localPath)
+      const normalizedText = text.trim()
+      const estimatedTokens = normalizedText ? estimateDocumentTokensConservatively(normalizedText) : 0
+      const canUseNativePdf = attachment.mediaType === 'application/pdf'
+        && normalizedText.length > 0
+        && attachment.size <= NATIVE_DOCUMENT_MAX_FILE_BYTES
+        && nativeBytes + attachment.size <= NATIVE_DOCUMENT_MAX_TOTAL_BYTES
+        && estimatedTokens <= NATIVE_DOCUMENT_MAX_TEXT_TOKENS
+        && prepared.nativeEstimatedTokens + estimatedTokens <= NATIVE_DOCUMENT_MAX_TOTAL_TEXT_TOKENS
+
+      if (canUseNativePdf) {
+        prepared.nativeDocumentIds.add(attachment.id)
+        prepared.nativeEstimatedTokens += estimatedTokens
+        nativeBytes += attachment.size
+      } else if (normalizedText) {
+        prepared.retrievalSources.push({ filename: attachment.filename, text: normalizedText })
+      } else {
+        prepared.notices.push(`[文档 ${attachment.filename} 内容为空]`)
+      }
+    } catch (error) {
+      const errorMsg = error instanceof Error ? error.message : '未知错误'
+      console.warn(`[聊天服务] 文档提取失败: ${attachment.filename}`, error)
+      prepared.notices.push(`[文档 ${attachment.filename} 内容提取失败: ${errorMsg}]`)
+    }
+  }
+  return prepared
+}
+
 // ===== 文档附件文本提取 =====
 
 /**
@@ -71,38 +140,19 @@ function getImageAttachmentData(attachments?: FileAttachment[]): ImageAttachment
  */
 async function enrichMessageWithDocuments(
   messageText: string,
-  attachments?: FileAttachment[],
+  prepared: PreparedDocuments,
 ): Promise<string> {
-  if (!attachments || attachments.length === 0) return messageText
-
-  // 筛选出文档类附件（非图片）
-  const docAttachments = attachments.filter((att) => isDocumentAttachment(att.mediaType))
-  if (docAttachments.length === 0) return messageText
-
-  const parts: string[] = [messageText]
-
-  for (const att of docAttachments) {
-    try {
-      const text = await extractTextFromAttachment(att.localPath)
-      if (text.trim()) {
-        parts.push(`\n<file name="${att.filename}">\n${text}\n</file>`)
-      } else {
-        parts.push(`\n<file name="${att.filename}">\n[文件内容为空]\n</file>`)
-      }
-    } catch (error) {
-      const errorMsg = error instanceof Error ? error.message : '未知错误'
-      console.warn(`[聊天服务] 文档提取失败: ${att.filename}`, error)
-      parts.push(`\n<file name="${att.filename}">\n[文件内容提取失败: ${errorMsg}]\n</file>`)
-    }
-  }
-
-  return parts.join('')
+  const context = retrieveDocumentContexts(prepared.retrievalSources, messageText)
+  const sections = [messageText]
+  if (context.content) sections.push(context.content)
+  if (prepared.notices.length > 0) sections.push(prepared.notices.join('\n'))
+  return sections.filter(Boolean).join('\n\n')
 }
 
 /**
  * 为历史消息列表注入文档附件文本
  *
- * 遍历历史消息，对包含文档附件的用户消息进行文本增强。
+ * 历史附件只保留引用，禁止每一轮重新注入整篇文档。
  * 返回新的消息数组（不修改原始消息）。
  */
 async function enrichHistoryWithDocuments(
@@ -115,8 +165,8 @@ async function enrichHistoryWithDocuments(
     if (msg.role === 'user' && msg.attachments && msg.attachments.length > 0) {
       const hasDocuments = msg.attachments.some((att) => isDocumentAttachment(att.mediaType))
       if (hasDocuments) {
-        const enrichedContent = await enrichMessageWithDocuments(msg.content, msg.attachments)
-        enriched.push({ ...msg, content: enrichedContent })
+        const names = msg.attachments.filter((att) => isDocumentAttachment(att.mediaType)).map((att) => att.filename)
+        enriched.push({ ...msg, content: `${msg.content}\n[历史附件引用：${names.join('、')}]` })
         continue
       }
     }
@@ -250,7 +300,21 @@ export async function sendMessage(
   // 5. 过滤历史并提取文档附件文本
   const filteredHistory = filterHistory(fullHistory, contextDividers, contextLength)
   const enrichedHistory = await enrichHistoryWithDocuments(filteredHistory)
-  const enrichedUserMessage = await enrichMessageWithDocuments(userMessage, attachments)
+  const preparedDocuments = await prepareDocuments(attachments)
+  const enrichedUserMessage = await enrichMessageWithDocuments(userMessage, preparedDocuments)
+  const estimatedPromptTokens = estimatePromptTokens([
+    systemMessage ?? '',
+    ...enrichedHistory.map((message) => message.content),
+    enrichedUserMessage,
+  ]) + preparedDocuments.nativeEstimatedTokens
+  // 留出模型回复和工具调用空间。超限必须由客户端明确提示，不能让 A2 静默截断。
+  if (estimatedPromptTokens > PROMPT_SAFE_INPUT_TOKENS) {
+    webContents.send(CHAT_IPC_CHANNELS.STREAM_ERROR, {
+      conversationId,
+      error: `当前对话预计 ${estimatedPromptTokens.toLocaleString()} tokens，超过快速稳定模式的安全输入预算 ${PROMPT_SAFE_INPUT_TOKENS.toLocaleString()}。请新建会话、缩小文档范围或先生成文档摘要。`,
+    })
+    return
+  }
 
   // 6. 创建 AbortController
   const controller = new AbortController()
@@ -330,6 +394,10 @@ export async function sendMessage(
         systemMessage: effectiveSystemMessage,
         attachments,
         readImageAttachments: getImageAttachmentData,
+        readDocumentAttachments: (requestAttachments) => getDocumentAttachmentData(
+          requestAttachments,
+          preparedDocuments.nativeDocumentIds,
+        ),
         thinkingEnabled,
         tools,
         continuationMessages: continuationMessages.length > 0 ? continuationMessages : undefined,
@@ -406,6 +474,10 @@ export async function sendMessage(
         systemMessage: effectiveSystemMessage,
         attachments,
         readImageAttachments: getImageAttachmentData,
+        readDocumentAttachments: (requestAttachments) => getDocumentAttachmentData(
+          requestAttachments,
+          preparedDocuments.nativeDocumentIds,
+        ),
         thinkingEnabled,
         // 不传 tools，强制模型生成文本回复而非继续调用工具
         continuationMessages,

apps/electron/src/main/lib/dingtalk-config.test.ts

@@ -3,6 +3,7 @@ import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync
 import * as os from 'node:os'
 import { join } from 'node:path'
 import { createStableDingTalkBotId } from './dingtalk-bot-identity'
+import { createElectronMock } from './test-helpers/electron-mock'
 
 type DingTalkConfigModule = typeof import('./dingtalk-config')
 type ConfigPathsModule = typeof import('./config-paths')
@@ -13,17 +14,9 @@ let tempHome: string
 const originalHome = process.env.HOME
 const originalPromaDev = process.env.PROMA_DEV
 
-mock.module('electron', () => ({
-  app: {
-    isPackaged: true,
-    getPath: () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'),
-  },
-  safeStorage: {
-    isEncryptionAvailable: () => false,
-    encryptString: (value: string) => Buffer.from(value),
-    decryptString: (value: Buffer) => value.toString('utf-8'),
-  },
-}))
+mock.module('electron', () => createElectronMock(
+  () => join(process.env.HOME ?? tempHome, 'Library', 'Application Support'),
+))
 
 mock.module('node:os', () => ({
   ...os,

apps/electron/src/main/lib/document-context.test.ts

@@ -0,0 +1,55 @@
+import { describe, expect, test } from 'bun:test'
+import {
+  chunkDocument,
+  DOCUMENT_CONTEXT_MAX_TOKENS,
+  estimatePromptTokens,
+  retrieveDocumentContexts,
+} from './document-context'
+
+describe('文档上下文预算', () => {
+  test('长文档按保守 token 预算重叠切分并检索中文相关内容', () => {
+    const text = `${'普通内容 '.repeat(4_000)}关键结论：财政政策有效。${'尾部 '.repeat(4_000)}`
+    const chunks = chunkDocument(text)
+    const result = retrieveDocumentContexts([{ filename: '财政报告.pdf', text }], '财政政策')
+
+    expect(chunks.length).toBeGreaterThan(1)
+    expect(chunks.every((chunk) => chunk.tokens <= 1_100)).toBeTrue()
+    expect(result.content).toContain('财政政策有效')
+  })
+
+  test('多份文档共享六万 token 总预算', () => {
+    const result = retrieveDocumentContexts([
+      { filename: '甲.pdf', text: `甲文档 财政政策 ${'研究资料 '.repeat(20_000)}` },
+      { filename: '乙.pdf', text: `乙文档 货币政策 ${'研究资料 '.repeat(20_000)}` },
+    ], '比较财政政策与货币政策')
+
+    expect(result.includedDocuments).toEqual(['甲.pdf', '乙.pdf'])
+    expect(result.estimatedTokens).toBeLessThanOrEqual(DOCUMENT_CONTEXT_MAX_TOKENS)
+  })
+
+  test('相同文档只注入一次', () => {
+    const text = `重复文档 ${'内容 '.repeat(2_000)}`
+    const result = retrieveDocumentContexts([
+      { filename: '原件.pdf', text },
+      { filename: '副本.pdf', text },
+    ], '重复文档')
+
+    expect(result.duplicateDocumentsSkipped).toBe(1)
+    expect(result.includedDocuments).toEqual(['原件.pdf'])
+  })
+
+  test('使用 Anthropic tokenizer 估算完整 prompt', () => {
+    expect(estimatePromptTokens(['hello world', '财政政策'])).toBeGreaterThan(0)
+  })
+
+  test('百万字符文档分块保持线性性能且不调用全文 tokenizer', () => {
+    const text = '财政政策与货币政策。'.repeat(100_000)
+    const startedAt = performance.now()
+    const chunks = chunkDocument(text)
+    const elapsedMs = performance.now() - startedAt
+
+    expect(chunks.length).toBeGreaterThan(1_000)
+    expect(chunks.every((chunk) => chunk.tokens <= 1_100)).toBeTrue()
+    expect(elapsedMs).toBeLessThan(5_000)
+  })
+})