From 12485a8ef5ceebbacf0997b31086751ea4fe9bab Mon Sep 17 00:00:00 2001 From: mrsibe Date: Tue, 29 Sep 2026 17:22:31 +0800 Subject: [PATCH 1/4] fix(embedding): run local inference in a worker and index imports in the background (#176) Local ONNX inference ran in the Electron main process, so a book-sized import held the event loop for the whole tokenize + forward pass and the window stalled. Ingestion was also a single long IPC call that resolved only after parse, chunk, embed and persist. - add embeddingWorker (worker_threads) with batching, per-batch progress and per-batch retry inside the worker - add WorkerEmbeddingBackend; EmbeddingService selects it in production and keeps the in-process backend for the injected test loader - build the worker as a second main-process entry (electron.vite.config.ts) - add IngestionQueue and route knowledge:add-files / add-folder through it: documents are registered as `pending` and returned immediately, indexing continues over knowledge:index-progress - broadcast index progress to all windows instead of the original sender - reload the document list when a background import completes or fails --- electron.vite.config.ts | 8 + src/main/embedding/WorkerEmbeddingBackend.ts | 203 ++++++++++++++++++ src/main/embedding/embeddingWorker.ts | 121 +++++++++++ src/main/embedding/embeddingWorkerProtocol.ts | 79 +++++++ src/main/embedding/types.ts | 6 + src/main/ipc/knowledgeHandlers.ts | 36 ++-- src/main/services/EmbeddingService.ts | 36 +++- src/main/services/KnowledgeService.ts | 138 ++++++++---- src/main/services/ingestion/IngestionQueue.ts | 74 +++++++ src/main/services/ingestion/folderWatch.ts | 3 +- .../notebook/source/DocumentList.tsx | 2 +- src/renderer/src/store/knowledgeStore.ts | 18 +- test/embeddingWorkerBackend.test.ts | 74 +++++++ test/fixtures/embeddingWorker.fake.mjs | 65 ++++++ test/ingestionQueue.test.ts | 85 ++++++++ 15 files changed, 880 insertions(+), 68 deletions(-) create mode 100644 src/main/embedding/WorkerEmbeddingBackend.ts create mode 100644 src/main/embedding/embeddingWorker.ts create mode 100644 src/main/embedding/embeddingWorkerProtocol.ts create mode 100644 src/main/services/ingestion/IngestionQueue.ts create mode 100644 test/embeddingWorkerBackend.test.ts create mode 100644 test/fixtures/embeddingWorker.fake.mjs create mode 100644 test/ingestionQueue.test.ts diff --git a/electron.vite.config.ts b/electron.vite.config.ts index 5efae6d..ae21071 100644 --- a/electron.vite.config.ts +++ b/electron.vite.config.ts @@ -45,6 +45,14 @@ export default defineConfig({ plugins: [copyMigrationsPlugin()], build: { rollupOptions: { + // Two entries: the Electron main process and the `worker_threads` entry that + // runs local ONNX inference off the main process (#176). Both land in + // out/main, so `WorkerEmbeddingBackend` can resolve the worker as a sibling + // of the main bundle. + input: { + index: resolve('src/main/index.ts'), + embeddingWorker: resolve('src/main/embedding/embeddingWorker.ts') + }, // Only things rollup CANNOT inline may stay external. Every entry in // package.json `dependencies` is already externalized automatically // (electron-vite's build.externalizeDeps defaults to true); the two diff --git a/src/main/embedding/WorkerEmbeddingBackend.ts b/src/main/embedding/WorkerEmbeddingBackend.ts new file mode 100644 index 0000000..26a2bcc --- /dev/null +++ b/src/main/embedding/WorkerEmbeddingBackend.ts @@ -0,0 +1,203 @@ +/** + * WorkerEmbeddingBackend（#176） + * + * `EmbeddingBackend` 的本地实现，但推理不在主线程：它把文本交给 `embeddingWorker` + * 线程，ONNX 的 tokenize + 前向在那里跑。主进程事件循环只剩消息收发。 + * + * 批次与进度由 worker 负责（见 embeddingWorkerProtocol.ts），所以这个后端声明 + * `batchesInternally = true`，EmbeddingService 不会在外面再切一遍。 + */ + +import { Worker } from 'worker_threads' +import { join } from 'path' +import type { EmbeddingPurpose, EmbeddingSpace } from '../../shared/types' +import Logger from '../../shared/utils/logger' +import { getLocalSpace } from './space' +import type { BackendEmbeddingResult, EmbeddingBackend } from './types' +import type { + EmbeddingWorkerInit, + EmbeddingWorkerResponse, + EmbeddingWorkerSnapshot +} from './embeddingWorkerProtocol' + +export interface WorkerEmbeddingBackendOptions { + /** userData/models */ + cacheDir: string + /** 一次 ONNX 前向的文本数 */ + batchSize: number + /** 单个 batch 失败后的重试上限（worker 内部执行） */ + maxRetries?: number + /** 重试基础退避（毫秒） */ + retryDelay?: number + /** 覆盖 worker 入口，测试注入用 */ + workerPath?: string + /** pipeline 首次加载进度（0-1） */ + onLoadProgress?: (progress: number) => void +} + +interface PendingRequest { + resolve: (results: BackendEmbeddingResult[]) => void + reject: (error: Error) => void + onProgress?: (completed: number, total: number) => void +} + +export class WorkerEmbeddingBackend implements EmbeddingBackend { + readonly kind = 'local' as const + /** 批处理在 worker 内完成，外层不要再切 */ + readonly batchesInternally = true + + private readonly cacheDir: string + private readonly batchSize: number + private readonly maxRetries: number + private readonly retryDelay: number + private readonly explicitWorkerPath?: string + private readonly onLoadProgress?: (progress: number) => void + private readonly space: EmbeddingSpace = getLocalSpace() + + private worker: Worker | null = null + private nextRequestId = 1 + private readonly pending = new Map() + + constructor(options: WorkerEmbeddingBackendOptions) { + this.cacheDir = options.cacheDir + this.batchSize = Math.max(1, options.batchSize) + this.maxRetries = Math.max(1, options.maxRetries ?? 3) + this.retryDelay = options.retryDelay ?? 1000 + this.explicitWorkerPath = options.workerPath + this.onLoadProgress = options.onLoadProgress + } + + getSpace(): EmbeddingSpace { + return this.space + } + + async isReady(): Promise { + return this.worker !== null + } + + async embed(text: string, purpose: EmbeddingPurpose): Promise { + const [result] = await this.embedBatch([text], purpose) + return result + } + + async embedBatch( + texts: string[], + purpose: EmbeddingPurpose, + onProgress?: (completed: number, total: number) => void + ): Promise { + if (texts.length === 0) { + return [] + } + + const worker = this.ensureWorker() + const id = this.nextRequestId++ + + return await new Promise((resolve, reject) => { + this.pending.set(id, { resolve, reject, onProgress }) + worker.postMessage({ type: 'embed', id, texts, purpose }) + }) + } + + /** + * 结束 worker 线程并释放 ONNX session。之后再调用会重新拉起一个 worker。 + */ + async dispose(): Promise { + const worker = this.worker + this.worker = null + this.failPending(new Error('Embedding worker disposed')) + + if (!worker) return + + const disposed = new Promise((resolve) => { + worker.once('exit', () => resolve()) + worker.once('error', () => resolve()) + }) + try { + worker.postMessage({ type: 'dispose', id: this.nextRequestId++ }) + } catch { + // worker 已经不在了，terminate 兜底 + } + worker.removeAllListeners('message') + await Promise.race([disposed, worker.terminate().then(() => undefined)]) + } + + private resolveWorkerPath(): string { + if (this.explicitWorkerPath) { + return this.explicitWorkerPath + } + // 主进程 bundle 与 worker 产物同在 out/main（electron.vite.config.ts 的第二个入口）。 + return join(__dirname, 'embeddingWorker.js') + } + + private ensureWorker(): Worker { + if (this.worker) { + return this.worker + } + + const init: EmbeddingWorkerInit = { + cacheDir: this.cacheDir, + batchSize: this.batchSize, + maxRetries: this.maxRetries, + retryDelay: this.retryDelay + } + const worker = new Worker(this.resolveWorkerPath(), { workerData: init }) + + worker.on('message', (message: EmbeddingWorkerResponse) => this.handleMessage(message)) + worker.on('error', (error) => { + Logger.error('WorkerEmbeddingBackend', 'Embedding worker failed:', error) + this.worker = null + this.failPending(error instanceof Error ? error : new Error(String(error))) + }) + worker.on('exit', (code) => { + if (this.worker === worker) { + this.worker = null + } + if (code !== 0) { + this.failPending(new Error(`Embedding worker exited with code ${code}`)) + } + }) + + this.worker = worker + return worker + } + + private handleMessage(message: EmbeddingWorkerResponse): void { + if (message.type === 'load-progress') { + this.onLoadProgress?.(message.progress) + return + } + + const pending = this.pending.get(message.id) + if (!pending) return + + if (message.type === 'progress') { + pending.onProgress?.(message.completed, message.total) + return + } + + this.pending.delete(message.id) + + if (message.type === 'result') { + pending.resolve(message.results.map(toBackendResult)) + } else if (message.type === 'error') { + pending.reject(new Error(message.message)) + } else if (message.type === 'disposed') { + // dispose 请求没有调用方等待；正常路径由 dispose() 自己结算 + } + } + + private failPending(error: Error): void { + for (const pending of this.pending.values()) { + pending.reject(error) + } + this.pending.clear() + } +} + +function toBackendResult(snapshot: EmbeddingWorkerSnapshot): BackendEmbeddingResult { + return { + embedding: snapshot.embedding, + model: snapshot.model, + dimensions: snapshot.dimensions + } +} diff --git a/src/main/embedding/embeddingWorker.ts b/src/main/embedding/embeddingWorker.ts new file mode 100644 index 0000000..2b4812a --- /dev/null +++ b/src/main/embedding/embeddingWorker.ts @@ -0,0 +1,121 @@ +/** + * embeddingWorker（#176） + * + * 本地 embedding 的 tokenize + ONNX 前向在 `worker_threads` 里跑，不在 Electron 主进程。 + * 主进程只负责发文本、收向量、转发进度；事件循环因此不再被一个 batch 的原生调用按住。 + * + * 批次与进度都在这里：worker 收到整批文本后自己按 `batchSize` 切开，每完成一批回一次 + * `progress`。这样「批大小」只有一处定义，主进程不会把一本书一次性推给 ONNX。 + */ + +import { parentPort, workerData } from 'worker_threads' +import { LocalEmbeddingBackend } from './LocalEmbeddingBackend' +import Logger from '../../shared/utils/logger' +import type { + EmbeddingWorkerInit, + EmbeddingWorkerRequest, + EmbeddingWorkerSnapshot +} from './embeddingWorkerProtocol' + +if (!parentPort) { + throw new Error('embeddingWorker must be started as a worker thread') +} + +/** 非空断言集中在这里：上面的检查之后，闭包里用到的就是这个常量。 */ +const port = parentPort + +const init = workerData as EmbeddingWorkerInit + +const backend = new LocalEmbeddingBackend({ + cacheDir: init.cacheDir, + onLoadProgress: (progress) => port.postMessage({ type: 'load-progress', progress }) +}) + +function chunk(items: T[], size: number): T[][] { + const batches: T[][] = [] + for (let index = 0; index < items.length; index += size) { + batches.push(items.slice(index, index + size)) + } + return batches +} + +async function handleEmbed( + id: number, + texts: string[], + purpose: 'query' | 'document' +): Promise { + const batches = chunk(texts, Math.max(1, init.batchSize)) + const results: EmbeddingWorkerSnapshot[] = [] + + for (const batch of batches) { + const batchResults = await embedBatchWithRetry(batch, purpose, id) + for (const result of batchResults) { + results.push({ + embedding: result.embedding, + model: result.model, + dimensions: result.dimensions + }) + } + port.postMessage({ + type: 'progress', + id, + completed: results.length, + total: texts.length + }) + } + + port.postMessage({ type: 'result', id, results }) +} + +/** + * 逐 batch 重试（而不是整批重来）：一次 book-sized 导入里某个 batch 抖一下，不应该让 + * 已经算好的几百条向量白做。重试策略跟 EmbeddingService 的远程路径同源。 + */ +async function embedBatchWithRetry( + batch: string[], + purpose: 'query' | 'document', + id: number +): Promise> { + let lastError: Error = new Error('Embedding failed') + + for (let attempt = 1; attempt <= Math.max(1, init.maxRetries); attempt++) { + try { + return await backend.embedBatch(batch, purpose) + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)) + Logger.warn( + 'EmbeddingWorker', + `attempt ${attempt} failed for request ${id}: ${lastError.message}` + ) + if (attempt < Math.max(1, init.maxRetries)) { + await new Promise((resolve) => setTimeout(resolve, init.retryDelay * 2 ** (attempt - 1))) + } + } + } + + throw lastError +} + +async function handle(message: EmbeddingWorkerRequest): Promise { + if (message.type === 'dispose') { + await backend.dispose() + port.postMessage({ type: 'disposed', id: message.id }) + return + } + + try { + await handleEmbed(message.id, message.texts, message.purpose) + } catch (error) { + port.postMessage({ + type: 'error', + id: message.id, + message: error instanceof Error ? error.message : String(error) + }) + } +} + +// 串行处理：ONNX session 不并发，排队也避免多个 inference 同时抢内存。 +let queue: Promise = Promise.resolve() +port.on('message', (message: EmbeddingWorkerRequest) => { + queue = queue.then(() => handle(message)).catch(() => undefined) +}) diff --git a/src/main/embedding/embeddingWorkerProtocol.ts b/src/main/embedding/embeddingWorkerProtocol.ts new file mode 100644 index 0000000..37b925f --- /dev/null +++ b/src/main/embedding/embeddingWorkerProtocol.ts @@ -0,0 +1,79 @@ +/** + * embeddingWorker 的线程间协议（#176）。 + * + * 主进程与服务层只通过这里定义的消息形状说话，worker 文件本身不导出实现细节。 + * 本地 ONNX 推理必须有这个边界：`LocalEmbeddingBackend` 的 tokenize + 前向是同步 + * 原生调用，跑在 Electron 主进程里会把事件循环按住，窗口拖拽、切页、输入全部跟着卡。 + */ + +import type { EmbeddingPurpose } from '../../shared/types' + +/** worker 启动参数（workerData） */ +export interface EmbeddingWorkerInit { + /** userData/models */ + cacheDir: string + /** 一次前向送进 pipeline 的文本数 */ + batchSize: number + /** 单个 batch 失败后的重试上限（与 EmbeddingService 的远程路径一致） */ + maxRetries: number + /** 重试基础退避（毫秒） */ + retryDelay: number +} + +/** 一条文本的向量结果（structured-clone 可传） */ +export interface EmbeddingWorkerSnapshot { + embedding: Float32Array + model: string + dimensions: number +} + +export interface EmbeddingWorkerEmbedRequest { + type: 'embed' + id: number + texts: string[] + purpose: EmbeddingPurpose +} + +export interface EmbeddingWorkerDisposeRequest { + type: 'dispose' + id: number +} + +export type EmbeddingWorkerRequest = EmbeddingWorkerEmbedRequest | EmbeddingWorkerDisposeRequest + +export interface EmbeddingWorkerProgressResponse { + type: 'progress' + id: number + completed: number + total: number +} + +/** pipeline 首次加载进度（0-1），与下载进度分开 */ +export interface EmbeddingWorkerLoadProgressResponse { + type: 'load-progress' + progress: number +} + +export interface EmbeddingWorkerResultResponse { + type: 'result' + id: number + results: EmbeddingWorkerSnapshot[] +} + +export interface EmbeddingWorkerErrorResponse { + type: 'error' + id: number + message: string +} + +export interface EmbeddingWorkerDisposedResponse { + type: 'disposed' + id: number +} + +export type EmbeddingWorkerResponse = + | EmbeddingWorkerProgressResponse + | EmbeddingWorkerLoadProgressResponse + | EmbeddingWorkerResultResponse + | EmbeddingWorkerErrorResponse + | EmbeddingWorkerDisposedResponse diff --git a/src/main/embedding/types.ts b/src/main/embedding/types.ts index 932c08d..5a669f4 100644 --- a/src/main/embedding/types.ts +++ b/src/main/embedding/types.ts @@ -38,4 +38,10 @@ export interface EmbeddingBackend { purpose: EmbeddingPurpose, onProgress?: (completed: number, total: number) => void ): Promise + + /** + * 后端自己在内部切 batch（#176 的 worker 后端）。缺省为 false：由 + * `EmbeddingService` 在外层分批，`embedBatch` 一次收到的就是一批。 + */ + readonly batchesInternally?: boolean } diff --git a/src/main/ipc/knowledgeHandlers.ts b/src/main/ipc/knowledgeHandlers.ts index 8b32cc0..4bb3c51 100644 --- a/src/main/ipc/knowledgeHandlers.ts +++ b/src/main/ipc/knowledgeHandlers.ts @@ -3,11 +3,31 @@ * 知识库相关的 IPC 处理函数 */ -import { ipcMain, IpcMainInvokeEvent, dialog, shell } from 'electron' +import { ipcMain, IpcMainInvokeEvent, dialog, shell, BrowserWindow } from 'electron' import { KnowledgeService } from '../services/KnowledgeService' import Logger from '../../shared/utils/logger' import { KnowledgeSchemas, validate } from './validation' +/** + * 后台索引进度广播（#176）。 + * + * 导入 IPC 现在在登记完 `pending` 行之后立即返回，进度来自后台队列；这时发起请求的 + * `event.sender` 可能已经销毁（换了窗口、关了标签），所以发给所有活着的窗口，而不是 + * 钉在当初那个 sender 上。 + */ +function broadcastIndexProgress(payload: { + notebookId?: string + documentId?: string + stage: string + progress: number +}): void { + BrowserWindow.getAllWindows().forEach((win) => { + if (!win.isDestroyed()) { + win.webContents.send('knowledge:index-progress', payload) + } + }) +} + /** * 注册知识库相关 IPC Handlers */ @@ -346,11 +366,7 @@ export function registerKnowledgeHandlers(knowledgeService: KnowledgeService) { params.notebookId, params.folderPath, (stage, progress) => { - event.sender.send('knowledge:index-progress', { - notebookId: params.notebookId, - stage, - progress - }) + broadcastIndexProgress({ notebookId: params.notebookId, stage, progress }) } ) @@ -397,15 +413,11 @@ export function registerKnowledgeHandlers(knowledgeService: KnowledgeService) { Logger.debug('KnowledgeHandlers', 'add-files:', params.paths.length) try { - const result = await knowledgeService.addDocumentsFromPaths( + const result = knowledgeService.enqueueDocumentsFromPaths( params.notebookId, params.paths, (stage, progress) => { - event.sender.send('knowledge:index-progress', { - notebookId: params.notebookId, - stage, - progress - }) + broadcastIndexProgress({ notebookId: params.notebookId, stage, progress }) } ) return { success: true, ...result } diff --git a/src/main/services/EmbeddingService.ts b/src/main/services/EmbeddingService.ts index 3657439..c59bb9e 100644 --- a/src/main/services/EmbeddingService.ts +++ b/src/main/services/EmbeddingService.ts @@ -20,6 +20,7 @@ import type { ConnectionManager } from '../models/ConnectionManager' import Logger from '../../shared/utils/logger' import { LocalEmbeddingBackend } from '../embedding/LocalEmbeddingBackend' import type { TransformersModuleLoader } from '../embedding/LocalEmbeddingBackend' +import { WorkerEmbeddingBackend } from '../embedding/WorkerEmbeddingBackend' import { RemoteEmbeddingBackend } from '../embedding/RemoteEmbeddingBackend' import type { BackendEmbeddingResult, EmbeddingBackend } from '../embedding/types' import { ModelDownloadService } from '../embedding/download/ModelDownloadService' @@ -49,10 +50,16 @@ export interface EmbeddingServiceOptions extends EmbeddingServiceConfig { getSources?: () => string[] | Promise /** 设置页里的下载进度广播 */ onDownloadProgress?: (progress: EmbeddingDownloadProgress) => void - /** 测试注入 */ + /** + * 测试注入。提供时本地推理退回进程内实现（假 pipeline），否则走 + * `worker_threads`（#176）。 + */ loadTransformers?: TransformersModuleLoader } +/** 本地后端都需要 Dispose（进程内是 ONNX session，worker 是线程）。 */ +type LocalBackend = EmbeddingBackend & { dispose(): Promise } + const DEFAULT_SOURCES = DEFAULT_EMBEDDING_SOURCES /** @@ -66,7 +73,7 @@ export class EmbeddingService { private readonly loadTransformers?: TransformersModuleLoader private config: Required - private localBackend: LocalEmbeddingBackend | null = null + private localBackend: LocalBackend | null = null private downloadService: ModelDownloadService | null = null private downloadProgressListener: ((progress: EmbeddingDownloadProgress) => void) | null = null private lastDetectedRemoteDimensions = 0 @@ -173,6 +180,13 @@ export class EmbeddingService { purpose: EmbeddingPurpose, onProgress?: (completed: number, total: number) => void ): Promise { + // worker 后端自己切 batch、逐批重试并在每次前向后回进度（#176）：重试与批次都 + // 只在 worker 里定义一次，这里不再用 withRetry 把整本书重跑一遍。 + if (backend.batchesInternally) { + Logger.info('EmbeddingService', `Local embedding via worker: ${texts.length} texts`) + return backend.embedBatch(texts, purpose, onProgress) + } + const batches = this.chunk(texts, this.config.localBatchSize) const results: BackendEmbeddingResult[] = [] @@ -313,12 +327,20 @@ export class EmbeddingService { this.config = { ...this.config, ...config } } - private getLocalBackend(): LocalEmbeddingBackend { + private getLocalBackend(): LocalBackend { if (!this.localBackend) { - this.localBackend = new LocalEmbeddingBackend({ - cacheDir: this.cacheDir, - loadTransformers: this.loadTransformers - }) + // 注入 loader 表示这是测试路径：没有真实 ONNX 可隔离，也就没有必要开线程。 + this.localBackend = this.loadTransformers + ? new LocalEmbeddingBackend({ + cacheDir: this.cacheDir, + loadTransformers: this.loadTransformers + }) + : new WorkerEmbeddingBackend({ + cacheDir: this.cacheDir, + batchSize: this.config.localBatchSize, + maxRetries: this.config.maxRetries, + retryDelay: this.config.retryDelay + }) } return this.localBackend } diff --git a/src/main/services/KnowledgeService.ts b/src/main/services/KnowledgeService.ts index 521b050..f628b07 100644 --- a/src/main/services/KnowledgeService.ts +++ b/src/main/services/KnowledgeService.ts @@ -70,6 +70,7 @@ import { type ReconcileResult } from './ingestion/folderWatch' import { scanFolder } from './ingestion/folderScan' +import { IngestionQueue } from './ingestion/IngestionQueue' import { deleteDocumentChunksFts, ensureChunksFts, indexChunksFts, searchChunksFts } from './fts' import { vectorStoreManager } from '../vectorstore' import Logger from '../../shared/utils/logger' @@ -192,6 +193,8 @@ export class KnowledgeService { private webFetchService: WebFetchService private knowledgeFilesDir: string private folderWatch: FolderWatchService + /** 后台索引队列（#176）：登记完就返回，解析/嵌入在这里继续。 */ + private readonly ingestionQueue: IngestionQueue constructor(embeddingService: EmbeddingService) { this.embeddingService = embeddingService @@ -204,6 +207,7 @@ export class KnowledgeService { this.ensureKnowledgeFilesDir() // 监听服务只持有本服务的引用，构造本身没有副作用（不会开数据库、不会开始监听）。 this.folderWatch = new FolderWatchService(this) + this.ingestionQueue = new IngestionQueue() } /** @@ -320,7 +324,10 @@ export class KnowledgeService { } /** - * 从文件添加文档 + * 从文件添加文档。 + * + * 这条路保持同步（smoke test / eval harness / 监听刷新需要它跑完再往下），后台 + * 队列只服务 `knowledge:add-files` 这类批量导入。 */ async addDocumentFromFile( notebookId: string, @@ -328,50 +335,23 @@ export class KnowledgeService { onProgress?: IndexProgressCallback, chunkOptions?: ChunkOptions ): Promise { - const db = getDatabase() - const documentId = `doc_${Date.now()}_${Math.random().toString(36).slice(2, 9)}` - const now = new Date() - - // source 行在**解析之前**就存在（#95）。以前解析失败时更新的是还没插入的行， - // 失败因此丢失，文档也无法重试。解析成功后再把解析结果回填。 - db.insert(documents) - .values({ - id: documentId, - notebookId, - title: basename(filePath) || 'Untitled', - type: 'file', - sourceUri: filePath, - status: 'processing', - chunkCount: 0, - createdAt: now, - updatedAt: now - }) - .run() - - // 先拷贝，再解析原文件：本地副本是重新索引/结构恢复时读取的东西。 - await this.ingestFile( - documentId, - filePath, - { copyFrom: filePath, chunkOptions }, - 'import', - onProgress - ) - + const documentId = this.createPendingFileDocument(notebookId, filePath) + await this.ingestPendingDocument(documentId, filePath, onProgress, chunkOptions) return documentId } /** - * 批量导入一组已经存在的文件路径（#98）。 + * 批量导入（#176）：同步登记，后台索引。 * - * 一个文件失败不会中止其余的：每个文件仍然走 `addDocumentFromFile`，所以每个都有自己 - * 的 ingestion run 与失败记录。已经在这个 notebook 里的路径被跳过并报告，而不是重复 - * 导入一份。 + * 返回时每个文件都已经是一个 `pending` 的 source；解析 / 分块 / 嵌入在 + * `IngestionQueue` 里继续，进度仍走 `knowledge:index-progress`。IPCC 不再把一次 + * book-sized 导入挂在调用栈上。 */ - async addDocumentsFromPaths( + enqueueDocumentsFromPaths( notebookId: string, paths: readonly string[], onProgress?: IndexProgressCallback - ): Promise { + ): BatchImportResult { const db = getDatabase() const existing = new Set( db @@ -387,35 +367,101 @@ export class KnowledgeService { const skipped: BatchImportSkip[] = [] const failed: BatchImportFailure[] = [] - for (let index = 0; index < paths.length; index++) { - const filePath = paths[index] - + for (const filePath of paths) { if (existing.has(filePath)) { skipped.push({ path: filePath, reason: 'already imported into this notebook' }) continue } try { - added.push(await this.addDocumentFromFile(notebookId, filePath)) + const documentId = this.createPendingFileDocument(notebookId, filePath) existing.add(filePath) + added.push(documentId) + this.ingestionQueue.enqueue({ + documentId, + onProgress, + run: async (jobProgress) => { + try { + await this.ingestPendingDocument(documentId, filePath, jobProgress) + } catch (error) { + // 失败也要给前台一个信号，否则后台跑挂的那一行会一直停在 pending。 + jobProgress('failed', 100) + throw error + } + } + }) } catch (error) { failed.push({ path: filePath, error: (error as Error).message }) } - - onProgress?.( - `importing ${index + 1}/${paths.length}`, - Math.round(((index + 1) / paths.length) * 100) - ) } return { added, skipped, failed } } + /** + * 登记一份 `pending` 的 source 行（不解析）。 + * + * source 行必须在解析之前存在（#95）：后台索引期间用户已经能在列表里看到它，也 + * 让失败有地方落。 + */ + private createPendingFileDocument(notebookId: string, filePath: string): string { + const documentId = `doc_${Date.now()}_${Math.random().toString(36).slice(2, 9)}` + const now = new Date() + + getDatabase() + .insert(documents) + .values({ + id: documentId, + notebookId, + title: basename(filePath) || 'Untitled', + type: 'file', + sourceUri: filePath, + status: 'pending', + chunkCount: 0, + createdAt: now, + updatedAt: now + }) + .run() + + return documentId + } + + /** 解析一份已登记的文件 source（同步与后台队列共用）。 */ + private async ingestPendingDocument( + documentId: string, + filePath: string, + onProgress?: IndexProgressCallback, + chunkOptions?: ChunkOptions + ): Promise { + const doc = this.getDocument(documentId) + if (!doc) { + Logger.warn('KnowledgeService', `Pending document ${documentId} disappeared before indexing`) + return + } + + getDatabase() + .update(documents) + .set({ status: 'processing', updatedAt: new Date() }) + .where(eq(documents.id, documentId)) + .run() + + // 先拷贝，再解析原文件：本地副本是重新索引/结构恢复时读取的东西。 + await this.ingestFile( + documentId, + filePath, + { copyFrom: filePath, chunkOptions }, + 'import', + onProgress + ) + } + /** * 把一个文件夹导入为**一次快照**（#98）。 * * 扫描只挑解析器认识的扩展名，其余文件被跳过并计数，而不是静默忽略。这是一次快照 * —— 之后新增到文件夹里的文件不会被自动带走，那是 #158 的监听。 + * + * 登记完成后立即返回，索引在后台队列里继续（#176）。 */ async addFolder( notebookId: string, @@ -423,7 +469,7 @@ export class KnowledgeService { onProgress?: IndexProgressCallback ): Promise { const scanned = await scanFolder(folderPath, this.fileParserService.supportedExtensions()) - return this.addDocumentsFromPaths( + return this.enqueueDocumentsFromPaths( notebookId, scanned.map((file) => file.path), onProgress diff --git a/src/main/services/ingestion/IngestionQueue.ts b/src/main/services/ingestion/IngestionQueue.ts new file mode 100644 index 0000000..e9513dc --- /dev/null +++ b/src/main/services/ingestion/IngestionQueue.ts @@ -0,0 +1,74 @@ +/** + * IngestionQueue（#176） + * + * 导入不再是一条长 IPC：handler 先把文档登记成 `pending` 并返回，解析 / 分块 / 嵌入 + * 在这里排队后台执行，进度经 `knowledge:index-progress` 继续推进。 + * + * 串行是有意的：本地 ONNX 只有一个 session，并发只会互相抢内存；而且一次导入的 + * 进度叙事只有一条线时用户才读得懂。 + */ + +import Logger from '../../../shared/utils/logger' + +/** 一个排队中的索引任务。`run` 收到的是这个任务的进度回调。 */ +export interface IngestionJob { + documentId: string + run: (onProgress: (stage: string, progress: number) => void) => Promise + onProgress?: (stage: string, progress: number) => void +} + +export class IngestionQueue { + private readonly jobs: IngestionJob[] = [] + /** 排队中或正在跑的 documentId；去重用，单看 jobs 会漏掉当前正在跑的那个。 */ + private readonly known = new Set() + private draining = false + private idleResolvers: Array<() => void> = [] + + /** 是否还有排队或正在跑的任务 */ + isBusy(): boolean { + return this.draining || this.jobs.length > 0 + } + + /** 排入一个任务；重复 documentId 会被忽略，避免同一份来源被索引两次。 */ + enqueue(job: IngestionJob): void { + if (this.known.has(job.documentId)) { + return + } + this.known.add(job.documentId) + this.jobs.push(job) + void this.drain() + } + + /** 队列清空后 resolve；测试与优雅退出用。 */ + whenIdle(): Promise { + if (!this.isBusy()) { + return Promise.resolve() + } + return new Promise((resolve) => this.idleResolvers.push(resolve)) + } + + private async drain(): Promise { + if (this.draining) return + this.draining = true + + try { + while (this.jobs.length > 0) { + const job = this.jobs.shift()! + try { + await job.run((stage, progress) => job.onProgress?.(stage, progress)) + } catch (error) { + // 失败已经写在 documents.status / ingestion run 上（#95）；队列继续跑下一个， + // 不能让一份坏文件把整批导入卡死。 + Logger.error('IngestionQueue', `Background indexing failed for ${job.documentId}:`, error) + } finally { + this.known.delete(job.documentId) + } + } + } finally { + this.draining = false + const resolvers = this.idleResolvers + this.idleResolvers = [] + for (const resolve of resolvers) resolve() + } + } +} diff --git a/src/main/services/ingestion/folderWatch.ts b/src/main/services/ingestion/folderWatch.ts index 3262c31..e5947df 100644 --- a/src/main/services/ingestion/folderWatch.ts +++ b/src/main/services/ingestion/folderWatch.ts @@ -140,7 +140,8 @@ export class FolderWatchService { } if (diff.added.length > 0) { - await this.knowledgeService.addDocumentsFromPaths( + // 后台队列（#176）：监听同步只负责登记，解析/嵌入在队列里继续。 + this.knowledgeService.enqueueDocumentsFromPaths( notebookId, diff.added.map((file) => file.path) ) diff --git a/src/renderer/src/components/notebook/source/DocumentList.tsx b/src/renderer/src/components/notebook/source/DocumentList.tsx index da70803..a6d136e 100644 --- a/src/renderer/src/components/notebook/source/DocumentList.tsx +++ b/src/renderer/src/components/notebook/source/DocumentList.tsx @@ -211,7 +211,7 @@ function DocumentItem({ {t('chunks', { count: document.chunkCount ?? 0 })} - {document.status === 'processing' && ( + {(document.status === 'processing' || document.status === 'pending') && ( {t('indexing')} diff --git a/src/renderer/src/store/knowledgeStore.ts b/src/renderer/src/store/knowledgeStore.ts index 31bb093..a63c3ab 100644 --- a/src/renderer/src/store/knowledgeStore.ts +++ b/src/renderer/src/store/knowledgeStore.ts @@ -29,6 +29,11 @@ interface KnowledgeStore { * #72 的失效引用判定只能在加载完成后做，否则冷启动的 deep-link 会在列表回来之前被误清。 */ documentsLoaded: boolean + /** + * 最近一次加载文档的 notebook。后台索引（#176）完成后只有进度事件可听， + * 靠它把刚转成 `indexed` 的那一行重新读回来。 + */ + activeNotebookId: string | null // Actions setDocuments: (docs: KnowledgeDocument[]) => void @@ -121,6 +126,7 @@ export const useKnowledgeStore = create()((set, get) => ({ indexProgress: null, error: null, documentsLoaded: false, + activeNotebookId: null, // Setters setDocuments: (documents) => set({ documents }), @@ -135,7 +141,7 @@ export const useKnowledgeStore = create()((set, get) => ({ // 加载文档列表 loadDocuments: async (notebookId) => { - set({ isLoading: true, error: null, documentsLoaded: false }) + set({ isLoading: true, error: null, documentsLoaded: false, activeNotebookId: notebookId }) try { const docs = await window.api.knowledge.getDocuments(notebookId) set({ documents: docs, isLoading: false, documentsLoaded: true }) @@ -314,6 +320,16 @@ export const useKnowledgeStore = create()((set, get) => ({ export function setupKnowledgeListeners(): () => void { const cleanupProgress = window.api.knowledge.onIndexProgress((data: IndexProgress) => { useKnowledgeStore.getState().setIndexProgress(data) + + // 后台导入的最后一步（#176）：这一行刚从 pending/processing 变成 indexed 或 failed。 + // 进度事件是唯一信号，所以在这里重读一次列表与统计。 + if (data.stage === 'completed' || data.stage === 'failed') { + const { activeNotebookId, loadDocuments, loadStats } = useKnowledgeStore.getState() + if (activeNotebookId) { + void loadDocuments(activeNotebookId) + void loadStats(activeNotebookId) + } + } }) return () => { diff --git a/test/embeddingWorkerBackend.test.ts b/test/embeddingWorkerBackend.test.ts new file mode 100644 index 0000000..a945207 --- /dev/null +++ b/test/embeddingWorkerBackend.test.ts @@ -0,0 +1,74 @@ +import { test } from 'node:test' +import assert from 'node:assert/strict' +import { fileURLToPath } from 'node:url' +import { dirname, join } from 'node:path' +import { WorkerEmbeddingBackend } from '../src/main/embedding/WorkerEmbeddingBackend.ts' + +/** + * #176：本地推理必须在 worker 线程里跑。 + * + * 这里用假 worker 钉住主进程侧的协议：整批文本交给 worker、批次在 worker 内切、进度按批 + * 回报、错误与线程崩溃都能传播给调用方，而不是把主进程挂死。 + */ + +const fixture = join( + dirname(fileURLToPath(import.meta.url)), + 'fixtures', + 'embeddingWorker.fake.mjs' +) + +test('worker 后端把整批文本交给 worker，并按批回报进度', async () => { + const service = new WorkerEmbeddingBackend({ + cacheDir: '/tmp/knownote-embedding-test', + batchSize: 16, + workerPath: fixture + }) + + const texts = Array.from({ length: 40 }, (_, index) => `chunk-${index}`) + const progress: Array<[number, number]> = [] + const results = await service.embedBatch(texts, 'document', (completed, total) => + progress.push([completed, total]) + ) + + assert.equal(results.length, 40) + // 批次由 worker 按 batchSize 切：进度是 16 → 32 → 40，而不是 40 一次跳到底。 + assert.deepEqual(progress, [ + [16, 40], + [32, 40], + [40, 40] + ]) + assert.ok(results[0].embedding instanceof Float32Array) + assert.equal(results[0].dimensions, 3) + + await service.dispose() +}) + +test('worker 的 inference 错误传播给调用方', async () => { + const service = new WorkerEmbeddingBackend({ + cacheDir: '/tmp/knownote-embedding-test', + batchSize: 2, + workerPath: fixture + }) + + await assert.rejects( + () => service.embedBatch(['fail-here'], 'document'), + /fake inference failure/ + ) + await service.dispose() +}) + +test('worker 崩溃时拒绝挂起的请求，且可以重新拉起', async () => { + const service = new WorkerEmbeddingBackend({ + cacheDir: '/tmp/knownote-embedding-test', + batchSize: 1, + workerPath: fixture + }) + + await assert.rejects(() => service.embedBatch(['crash-now'], 'document')) + + // 崩溃后 worker 被丢弃；下一次调用重新创建一个，而不是永远卡死。 + const results = await service.embedBatch(['ok'], 'document') + assert.equal(results.length, 1) + + await service.dispose() +}) diff --git a/test/fixtures/embeddingWorker.fake.mjs b/test/fixtures/embeddingWorker.fake.mjs new file mode 100644 index 0000000..119f620 --- /dev/null +++ b/test/fixtures/embeddingWorker.fake.mjs @@ -0,0 +1,65 @@ +/** + * embeddingWorker 协议的假 worker（#176 测试用）。 + * + * 真实 worker 会加载 transformers.js，那需要下载模型；这里只验证主进程侧的协议： + * 分批、进度、结果回填、错误传播与线程生命周期。故障由文本内容触发（`fail*` / `crash*`）， + * 这样不必往生产协议里塞测试开关。 + * + * 这是纯 JS 的测试替身，不参与类型检查，因此关掉需要标注返回类型的规则。 + */ +/* eslint-disable @typescript-eslint/explicit-function-return-type */ + +import { parentPort, workerData } from 'worker_threads' + +const port = parentPort +if (!port) { + throw new Error('fake embedding worker must run as a worker thread') +} + +const { batchSize = 16 } = workerData ?? {} + +function embed(texts, purpose) { + return texts.map((text) => { + const size = purpose === 'query' ? 2 : 3 + return { + embedding: new Float32Array(size).fill(text.length), + model: 'fake-model', + dimensions: size + } + }) +} + +port.on('message', (message) => { + if (message.type === 'dispose') { + port.postMessage({ type: 'disposed', id: message.id }) + port.close() + return + } + + if (message.texts.some((text) => text.startsWith('crash'))) { + process.exit(7) + return + } + + if (message.texts.some((text) => text.startsWith('fail'))) { + port.postMessage({ type: 'error', id: message.id, message: 'fake inference failure' }) + return + } + const batches = [] + for (let index = 0; index < message.texts.length; index += batchSize) { + batches.push(message.texts.slice(index, index + batchSize)) + } + + const results = [] + for (const batch of batches) { + results.push(...embed(batch, message.purpose)) + port.postMessage({ + type: 'progress', + id: message.id, + completed: results.length, + total: message.texts.length + }) + } + + port.postMessage({ type: 'result', id: message.id, results }) +}) diff --git a/test/ingestionQueue.test.ts b/test/ingestionQueue.test.ts new file mode 100644 index 0000000..9d246a0 --- /dev/null +++ b/test/ingestionQueue.test.ts @@ -0,0 +1,85 @@ +import { test } from 'node:test' +import assert from 'node:assert/strict' +import { IngestionQueue } from '../src/main/services/ingestion/IngestionQueue.ts' + +/** + * #176：批量导入的后台队列。 + * + * handler 登记完 `pending` 行就返回，解析/嵌入在这里排队。这些用例钉住队列的三条 + * 语义：串行、进度转发、单份失败不拖垮整批。 + */ + +const tick = () => new Promise((resolve) => setImmediate(resolve)) + +test('任务串行执行，且进度转发给入队方', async () => { + const queue = new IngestionQueue() + const order: string[] = [] + const progress: Array<[string, string, number]> = [] + + queue.enqueue({ + documentId: 'a', + onProgress: (stage, value) => progress.push(['a', stage, value]), + run: async (onProgress) => { + order.push('a:start') + onProgress('parsing', 5) + await tick() + order.push('a:end') + } + }) + queue.enqueue({ + documentId: 'b', + onProgress: (stage, value) => progress.push(['b', stage, value]), + run: async () => { + order.push('b:start') + await tick() + order.push('b:end') + } + }) + + await queue.whenIdle() + + // a 必须整体结束之后 b 才开始，不能交错。 + assert.deepEqual(order, ['a:start', 'a:end', 'b:start', 'b:end']) + assert.deepEqual(progress, [['a', 'parsing', 5]]) + assert.equal(queue.isBusy(), false) +}) + +test('一份失败不中止后面的任务', async () => { + const queue = new IngestionQueue() + const completed: string[] = [] + + queue.enqueue({ documentId: 'bad', run: async () => Promise.reject(new Error('boom')) }) + queue.enqueue({ + documentId: 'good', + run: async () => { + completed.push('good') + } + }) + + await queue.whenIdle() + assert.deepEqual(completed, ['good']) +}) + +test('同一个 documentId 不会重复入队', async () => { + const queue = new IngestionQueue() + let runs = 0 + + const job = { + documentId: 'same', + run: async () => { + runs += 1 + await tick() + } + } + queue.enqueue(job) + queue.enqueue(job) + + await queue.whenIdle() + assert.equal(runs, 1) +}) + +test('whenIdle 在空队列上立即 resolve', async () => { + const queue = new IngestionQueue() + await queue.whenIdle() + assert.equal(queue.isBusy(), false) +}) From 3de1d687c6b6c01143ebd937488dc2eb8c97c400 Mon Sep 17 00:00:00 2001 From: mrsibe Date: Tue, 29 Sep 2026 17:23:48 +0800 Subject: [PATCH 2/4] perf(chat): coalesce streaming snapshots and render each message from its own turn (#177) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A streaming answer wrote to the store on every snapshot, remapped the whole message array, and re-rendered every historical MessageItem — each with a Markdown parse, syntax highlight and KaTeX pass. - hold the latest snapshot per message and commit on a 40ms timer; flush synchronously before an outcome so terminal state is never delayed - keep the per-message sequence counter out of the store and write it only when a gap is actually detected - guard the pending→streaming transition so a chunk that changes nothing does not commit - render live content and reasoning from the message's own turn, and memoize MessageItem so unchanged answers do not re-render --- .../components/notebook/chat/MessageItem.tsx | 36 +++-- src/renderer/src/store/chatStore.ts | 140 ++++++++++++----- test/chatSnapshotThrottle.test.ts | 143 ++++++++++++++++++ test/chatTurnEvents.test.ts | 13 +- 4 files changed, 276 insertions(+), 56 deletions(-) create mode 100644 test/chatSnapshotThrottle.test.ts diff --git a/src/renderer/src/components/notebook/chat/MessageItem.tsx b/src/renderer/src/components/notebook/chat/MessageItem.tsx index f612939..2006d8e 100644 --- a/src/renderer/src/components/notebook/chat/MessageItem.tsx +++ b/src/renderer/src/components/notebook/chat/MessageItem.tsx @@ -1,4 +1,4 @@ -import { ReactElement, useState } from 'react' +import { memo, ReactElement, useState } from 'react' import ReactMarkdown from 'react-markdown' import remarkGfm from 'remark-gfm' import remarkMath from 'remark-math' @@ -28,6 +28,7 @@ import { citationToSourceAnchor, sourceAnchorsEqual } from '../../../../../shared/utils/sourceAnchor' +import { messageReasoning, messageText } from '../../../../../shared/utils/uiMessage' import { answerNoticeKey, canRecover, isAnswerLive } from '../../../../../shared/utils/answerState' import { useChatStore } from '../../../store/chatStore' import { Button } from '../../ui/button' @@ -40,7 +41,7 @@ interface MessageItemProps { message: ChatMessage } -export default function MessageItem({ message }: MessageItemProps): ReactElement { +function MessageItem({ message }: MessageItemProps): ReactElement { const { t } = useTranslation(['common', 'chat']) const isUser = message.role === 'user' const isSystem = message.role === 'system' @@ -62,6 +63,12 @@ export default function MessageItem({ message }: MessageItemProps): ReactElement // Whether the model is still thinking. The one live detail that is not on the // message: it comes from the assembled answer, and only exists while one does. const reasoningLive = useChatStore((state) => state.turns[message.id]?.reasoningLive ?? false) + // The live answer is read from the turn this message owns (#177), not from a copy + // remapped in `messages` on every chunk. Only this component subscribes to this + // slice, so while one answer streams the others do not re-render at all. + const liveMessage = useChatStore((state) => state.turns[message.id]?.message) + const content = liveMessage ? messageText(liveMessage) : message.content + const reasoningContent = liveMessage ? messageReasoning(liveMessage) : message.reasoningContent const retryMessage = useChatStore((state) => state.retryMessage) const continueMessage = useChatStore((state) => state.continueMessage) @@ -79,7 +86,7 @@ export default function MessageItem({ message }: MessageItemProps): ReactElement // Citation coverage (#156): which sentences carry a resolvable citation. Computed // from the answer and the persisted citations with the same pure function the // tests exercise, so "grounded" means one thing in the UI and in a test. - const coverage = message.content ? classifyClaimSupport(message.content, citations) : null + const coverage = content ? classifyClaimSupport(content, citations) : null // A chip is disabled only once the library has been read and the document is // provably absent. Before that (or after a failed load) the list is "unknown", @@ -104,7 +111,7 @@ export default function MessageItem({ message }: MessageItemProps): ReactElement const handleCopy = async () => { try { - await navigator.clipboard.writeText(message.content) + await navigator.clipboard.writeText(content) setCopied(true) setTimeout(() => setCopied(false), 2000) } catch (error) { @@ -117,7 +124,7 @@ export default function MessageItem({ message }: MessageItemProps): ReactElement if (!currentNotebook) return try { - await createNote(currentNotebook.id, message.content) + await createNote(currentNotebook.id, content) setAddedToNote(true) setTimeout(() => setAddedToNote(false), 2000) } catch (error) { @@ -212,11 +219,11 @@ export default function MessageItem({ message }: MessageItemProps): ReactElement
{/* Reasoning process display - only shown when reasoning content exists */} - {message.reasoningContent && ( - + {reasoningContent && ( + )} - {message.content ? ( + {content ? (
- {message.content} + {content} {/* Streaming message cursor */} {isLive && ( @@ -275,7 +282,7 @@ export default function MessageItem({ message }: MessageItemProps): ReactElement {/* What the answer was built from. Only after the turn ends: during streaming there is nothing to show yet, and an empty evidence list mid-answer would read as "nothing was used". */} - {message.content && !isLive && ( + {content && !isLive && ( )} {/* Action buttons - only shown when reply is complete and has content */} - {message.content && !isLive && ( + {content && !isLive && (
{/* Recovery, when the answer stopped early and asking again may help */} {canRecover(message.status) && ( @@ -389,3 +396,10 @@ export default function MessageItem({ message }: MessageItemProps): ReactElement
) } + +/** + * Memoized on the message's own record (#177). While one answer streams, every + * historical `MessageItem` keeps the same `message` object, so React skips it + * entirely — no Markdown parse, no syntax highlight, no KaTeX pass. + */ +export default memo(MessageItem) diff --git a/src/renderer/src/store/chatStore.ts b/src/renderer/src/store/chatStore.ts index 0eb30e6..ed16400 100644 --- a/src/renderer/src/store/chatStore.ts +++ b/src/renderer/src/store/chatStore.ts @@ -102,6 +102,58 @@ interface TurnAssembler { */ const assemblers = new Map() +/** + * Snapshot coalescing (#177). + * + * The assembler emits one snapshot per chunk — a reasoning model can produce + * hundreds per second — and each one used to be its own store commit and its own + * full-transcript render. The latest snapshot per message is held here and applied + * on a short timer, so the commit rate is bounded while no snapshot is ever + * dropped: a later one replaces the pending one for the same message. + */ +const SNAPSHOT_FLUSH_MS = 40 +const pendingSnapshots = new Map() +let snapshotTimer: ReturnType | null = null + +/** + * The last sequence number per message, kept out of the store (#177). + * + * Writing it on every chunk was a store commit per chunk — the very cadence the + * snapshot throttle exists to remove. It is only observable state when there is a + * gap, so that is the only time it reaches the store. + */ +const lastSequences = new Map() + +/** Apply every held snapshot in one commit. Terminal paths call this first. */ +function flushPendingSnapshots(): void { + if (snapshotTimer !== null) { + clearTimeout(snapshotTimer) + snapshotTimer = null + } + if (pendingSnapshots.size === 0) return + + const pending = new Map(pendingSnapshots) + pendingSnapshots.clear() + + useChatStore.setState((state) => { + const turns = { ...state.turns } + for (const [messageId, message] of pending) { + const turn = turns[messageId] + if (!turn) continue + turns[messageId] = { ...turn, message, reasoningLive: isReasoningLive(message) } + } + return { turns } + }) +} + +function scheduleSnapshotFlush(): void { + if (snapshotTimer !== null) return + snapshotTimer = setTimeout(() => { + snapshotTimer = null + flushPendingSnapshots() + }, SNAPSHOT_FLUSH_MS) +} + /** * What the overlay knows before the first event of a turn arrives. * @@ -183,6 +235,10 @@ function feedAssembler(messageId: string, chunk: UIMessageChunk): void { * is supposed to keep. */ function applyOutcome(event: Extract): void { + // Terminal state is applied now, not on the next turn of the throttle: the last + // snapshot is what the turn's content must settle on. + flushPendingSnapshots() + const state = useChatStore.getState() const turn = state.turns[event.messageId] @@ -213,29 +269,15 @@ function applyOutcome(event: Extract): void if (turn) state.setStreamingMessage(turn.notebookId, null) assemblers.delete(event.messageId) + lastSequences.delete(event.messageId) } function applySnapshot(messageId: string, message: UIMessage): void { - const text = messageText(message) - const reasoningContent = messageReasoning(message) - const reasoningLive = isReasoningLive(message) - - useChatStore.setState((state) => { - const turn = state.turns[messageId] - const turns = turn - ? { ...state.turns, [messageId]: { ...turn, message, reasoningLive } } - : state.turns - - // Only touch the visible list when this is the message the reader is looking at; - // a turn in another notebook still assembles so switching back is instant. - const messages = state.messages.some((item) => item.id === messageId) - ? state.messages.map((item) => - item.id === messageId ? { ...item, content: text, reasoningContent } : item - ) - : state.messages - - return { turns, messages } - }) + // Hold the latest snapshot; the timer commits it together with any other live + // turn. This is the whole point of #177: one chunk must not be one store write + // and one render of every historical message. + pendingSnapshots.set(messageId, message) + scheduleSnapshotFlush() } export const useChatStore = create()((set, get) => ({ @@ -513,35 +555,49 @@ export function setupChatListeners(): () => void { const turn = state.turns[event.messageId] // A gap means content was lost on the way. Nothing replays yet, so say so - // rather than showing an answer that is quietly short. - if (turn && !isSequenceContinuing(turn.lastSeq, event.seq)) { - console.warn( - `[ChatStore] missing events for ${event.messageId}: expected ${(turn.lastSeq ?? 0) + 1}, got ${event.seq}` - ) - useChatStore.setState((current) => ({ - turns: { - ...current.turns, - [event.messageId]: { ...turn, lastSeq: event.seq, sequenceGap: true } - } - })) - } else if (turn) { - useChatStore.setState((current) => ({ - turns: { ...current.turns, [event.messageId]: { ...turn, lastSeq: event.seq } } - })) + // rather than showing an answer that is quietly short. The running sequence + // stays out of the store (#177): writing it per chunk was a commit per + // chunk, which is exactly the cadence the snapshot throttle removes. + if (turn) { + const previous = lastSequences.get(event.messageId) + if (!isSequenceContinuing(previous, event.seq)) { + console.warn( + `[ChatStore] missing events for ${event.messageId}: expected ${(previous ?? 0) + 1}, got ${event.seq}` + ) + useChatStore.setState((current) => ({ + turns: { + ...current.turns, + [event.messageId]: { + ...(current.turns[event.messageId] ?? turn), + lastSeq: event.seq, + sequenceGap: true + } + } + })) + } + lastSequences.set(event.messageId, event.seq) } feedAssembler(event.messageId, event.event) // The first event of a turn is what turns "waiting" into "streaming", the - // same transition the execution makes in Main. + // same transition the execution makes in Main. Guarded so a chunk that + // changes nothing does not commit: returning the current state is a no-op + // for zustand, and a per-chunk commit is what #177 removes. if (event.event.type !== 'start') { - useChatStore.setState((current) => ({ - messages: current.messages.map((message) => - message.id === event.messageId && message.status === 'pending' - ? { ...message, status: 'streaming' } - : message + useChatStore.setState((current) => { + const transition = current.messages.some( + (message) => message.id === event.messageId && message.status === 'pending' ) - })) + if (!transition) return current + return { + messages: current.messages.map((message) => + message.id === event.messageId && message.status === 'pending' + ? { ...message, status: 'streaming' } + : message + ) + } + }) } break } diff --git a/test/chatSnapshotThrottle.test.ts b/test/chatSnapshotThrottle.test.ts new file mode 100644 index 0000000..f98a58c --- /dev/null +++ b/test/chatSnapshotThrottle.test.ts @@ -0,0 +1,143 @@ +import { test } from 'node:test' +import assert from 'node:assert/strict' +import type { UIMessageChunk } from 'ai' +import { setupChatListeners, useChatStore } from '../src/renderer/src/store/chatStore.ts' +import type { ChatMessage, ChatTurnEvent } from '../src/shared/types/chat.ts' +import { messageText } from '../src/shared/utils/uiMessage.ts' + +/** + * #177: a streaming answer must not re-render the whole transcript per chunk. + * + * The assembler emits a snapshot per chunk; the store holds the latest one per + * message and commits on a short timer. These tests pin the three properties the + * acceptance names: coalescing, last-one-wins, and terminal state applied now — + * never on the next turn of the throttle. + */ + +let deliver: ((event: ChatTurnEvent) => void) | undefined + +function installApi(): void { + ;(globalThis as unknown as { window: unknown }).window = { + api: { + onTurnEvent: (callback: (event: ChatTurnEvent) => void) => { + deliver = callback + return () => { + deliver = undefined + } + }, + getActiveSession: async () => null + } + } +} + +const waitFor = async (condition: () => boolean, timeoutMs = 1000): Promise => { + const deadline = Date.now() + timeoutMs + while (!condition()) { + if (Date.now() > deadline) throw new Error('the condition never became true') + await new Promise((resolve) => setTimeout(resolve, 5)) + } +} + +const placeholder = (messageId: string): ChatMessage => ({ + id: messageId, + sessionId: 'session_1', + notebookId: 'notebook_1', + role: 'assistant', + content: '', + reasoningContent: null, + status: 'pending', + finishReason: null, + error: null, + usage: null, + finishedAt: null, + attemptOf: null, + createdAt: new Date() +}) + +function seedTurn(messageId: string): void { + useChatStore.setState({ + messages: [placeholder(messageId)], + turns: { [messageId]: { notebookId: 'notebook_1' } }, + streamingMessages: { notebook_1: messageId }, + currentSession: null, + sessions: [] + }) +} + +const chunk = (messageId: string, seq: number, event: UIMessageChunk): ChatTurnEvent => ({ + type: 'chunk', + executionId: 'exec_1', + messageId, + seq, + event +}) + +const liveTextOf = (messageId: string): string => + messageText(useChatStore.getState().turns[messageId]?.message) + +test('many chunks collapse into far fewer store commits', async () => { + installApi() + const close = setupChatListeners() + seedTurn('msg_throttle') + + let commits = 0 + const unsubscribe = useChatStore.subscribe(() => { + commits += 1 + }) + + deliver?.(chunk('msg_throttle', 1, { type: 'text-start', id: 't1' })) + for (let index = 0; index < 120; index++) { + deliver?.(chunk('msg_throttle', index + 2, { type: 'text-delta', id: 't1', delta: 'x' })) + } + + await waitFor(() => liveTextOf('msg_throttle').length === 120) + unsubscribe() + + // 120 chunks on one tick must not be 120 commits. The only writes are the + // pending→streaming transition and the flush itself. + assert.ok(commits <= 4, `expected ≤4 commits for 120 chunks, got ${commits}`) + assert.equal(liveTextOf('msg_throttle'), 'x'.repeat(120)) + + close() +}) + +test('a snapshot inside the throttle window is not dropped: last one wins', async () => { + installApi() + const close = setupChatListeners() + seedTurn('msg_last') + + deliver?.(chunk('msg_last', 1, { type: 'text-start', id: 't1' })) + deliver?.(chunk('msg_last', 2, { type: 'text-delta', id: 't1', delta: 'first' })) + deliver?.(chunk('msg_last', 3, { type: 'text-delta', id: 't1', delta: '-second' })) + + // Both deltas land inside one throttle window. If the pending snapshot were + // dropped rather than replaced, the text would read "first" or "-second". + await waitFor(() => liveTextOf('msg_last') === 'first-second') + + close() +}) + +test('the terminal outcome is applied synchronously, not on the timer', async () => { + installApi() + const close = setupChatListeners() + seedTurn('msg_terminal') + + deliver?.(chunk('msg_terminal', 1, { type: 'text-start', id: 't1' })) + deliver?.(chunk('msg_terminal', 2, { type: 'text-delta', id: 't1', delta: 'the answer' })) + deliver?.({ + type: 'outcome', + executionId: 'exec_1', + messageId: 'msg_terminal', + seq: 3, + outcome: { status: 'completed' }, + finishReason: 'stop', + messageMetadata: {} + }) + + // No timer wait: once the turn settles, the content is already there. + await waitFor(() => useChatStore.getState().streamingMessages['notebook_1'] === undefined) + const message = useChatStore.getState().messages.find((item) => item.id === 'msg_terminal') + assert.equal(message?.content, 'the answer') + + close() +}) diff --git a/test/chatTurnEvents.test.ts b/test/chatTurnEvents.test.ts index e9232d2..5e64798 100644 --- a/test/chatTurnEvents.test.ts +++ b/test/chatTurnEvents.test.ts @@ -8,6 +8,7 @@ import { isAnswerLive, keepsPartialAnswer } from '../src/shared/utils/answerState.ts' +import { messageText } from '../src/shared/utils/uiMessage.ts' /** * The renderer's half of the stream protocol (#141). @@ -105,6 +106,10 @@ const textStream = (messageId: string, text: string): ChatTurnEvent[] => [ const messageOf = (messageId: string): ChatMessage | undefined => useChatStore.getState().messages.find((message) => message.id === messageId) +/** The live answer as the UI reads it now (#177): from the turn, not the list. */ +const liveTextOf = (messageId: string): string => + messageText(useChatStore.getState().turns[messageId]?.message) + test('the events assemble into the answer on screen', async () => { installApi() const close = setupChatListeners() @@ -112,7 +117,9 @@ test('the events assemble into the answer on screen', async () => { for (const event of textStream('msg_assemble', 'the answer')) deliver?.(event) - await waitFor(() => messageOf('msg_assemble')?.content === 'the answer') + // The streaming answer lives on the turn (#177); the message list is settled on + // the outcome so a chunk never remaps it. + await waitFor(() => liveTextOf('msg_assemble') === 'the answer') // The first event is what turns "waiting" into "streaming", the same transition // the execution makes in Main — and only the outcome event ends the turn, which is @@ -143,7 +150,7 @@ test('the outcome ends the turn and applies what it carried', async () => { seedTurn('msg_outcome') for (const event of textStream('msg_outcome', 'the answer')) deliver?.(event) - await waitFor(() => messageOf('msg_outcome')?.content === 'the answer') + await waitFor(() => liveTextOf('msg_outcome') === 'the answer') deliver?.({ type: 'outcome', @@ -243,7 +250,7 @@ test('a continuation appends to the answer on screen instead of replacing it', a deliver?.(chunk('msg_continue', 2, { type: 'text-delta', id: 't1', delta: 'and the rest' })) deliver?.(chunk('msg_continue', 3, { type: 'text-end', id: 't1' })) - await waitFor(() => messageOf('msg_continue')?.content === 'the beginning. and the rest') + await waitFor(() => liveTextOf('msg_continue') === 'the beginning. and the rest') close() }) From c6691b46bf632892f9ad1692d656bdb562d897b8 Mon Sep 17 00:00:00 2001 From: mrsibe Date: Tue, 29 Sep 2026 17:24:23 +0800 Subject: [PATCH 3/4] fix(toast): mount one root toaster and key repeatable toasts (#178) lived inside NoteEditor, and sonner replays every still-active toast to a late subscriber. Import, save and excerpt toasts fired while no note was open were therefore held and released together the next time a note opened. - mount a single at the app root and remove it from NoteEditor - give repeatable toasts a stable id so a later call updates the existing toast instead of stacking a duplicate - add a guard test that pins the single mount and the keyed call sites --- src/renderer/src/App.tsx | 8 ++ .../src/components/notebook/NotePanel.tsx | 4 +- .../src/components/notebook/SourcePanel.tsx | 12 ++- .../components/notebook/note/NoteEditor.tsx | 4 +- test/toastRoot.test.ts | 84 +++++++++++++++++++ 5 files changed, 105 insertions(+), 7 deletions(-) create mode 100644 test/toastRoot.test.ts diff --git a/src/renderer/src/App.tsx b/src/renderer/src/App.tsx index 5b76c1f..c40e63c 100644 --- a/src/renderer/src/App.tsx +++ b/src/renderer/src/App.tsx @@ -9,6 +9,7 @@ import MindMapPage from './components/pages/MindMapPage' import QuizPage from './components/pages/QuizPage' import AnkiPage from './components/pages/AnkiPage' import SettingsDialog from './components/settings/SettingsDialog' +import { Toaster } from './components/ui/sonner' import { setupChatListeners } from './store/chatStore' import { useThemeStore } from './store/themeStore' import { useNotebookStore } from './store/notebookStore' @@ -107,6 +108,13 @@ function App(): React.JSX.Element { } /> + {/* + * One toaster for the whole app, mounted here rather than inside a note editor + * (#178). sonner replays every still-active toast to a subscriber that arrives + * late, so a toaster that only existed while a note was open held import, + * save and excerpt toasts and released them all at once on the next open. + */} + ) diff --git a/src/renderer/src/components/notebook/NotePanel.tsx b/src/renderer/src/components/notebook/NotePanel.tsx index 99ea564..65e35f5 100644 --- a/src/renderer/src/components/notebook/NotePanel.tsx +++ b/src/renderer/src/components/notebook/NotePanel.tsx @@ -208,7 +208,9 @@ export default function NotePanel(): ReactElement { const handleSave = async (title: string, content: string) => { if (!currentNote) return await updateNote(currentNote.id, { title, content }) - toast.success(t('noteSaved')) + // One toast per note (#178): saving the same note again updates it rather than + // stacking another "saved" line. + toast.success(t('noteSaved'), { id: `note-saved:${currentNote.id}` }) } // 删除笔记 diff --git a/src/renderer/src/components/notebook/SourcePanel.tsx b/src/renderer/src/components/notebook/SourcePanel.tsx index fe38f16..f5dc522 100644 --- a/src/renderer/src/components/notebook/SourcePanel.tsx +++ b/src/renderer/src/components/notebook/SourcePanel.tsx @@ -444,7 +444,11 @@ export default function SourcePanel(): ReactElement { if (currentNote) { requestAppendExcerpt(currentNote.id, excerptMarkdown) - toast.success(t('excerptAppended', { title: currentNote.title })) + // Stable id (#178): appending twice to the same note updates one toast + // instead of stacking a second identical one. + toast.success(t('excerptAppended', { title: currentNote.title }), { + id: `excerpt:${currentNote.id}` + }) return } @@ -452,7 +456,7 @@ export default function SourcePanel(): ReactElement { await createNote(notebookId, excerptMarkdown, t('excerptNoteTitle', { title: sourceTitle })) } catch (error) { console.error('[SourcePanel] Failed to save the excerpt as a note:', error) - toast.error(t('excerptSaveFailed')) + toast.error(t('excerptSaveFailed'), { id: 'excerpt-save' }) } }, [notebookId, currentNote, createNote, t] @@ -469,7 +473,9 @@ export default function SourcePanel(): ReactElement { const reportImportFailure = useCallback( (name: string, result: { success: boolean; error?: string }, reason?: string): void => { if (result.success) return - toast.error(t('importFailed', { name, error: reason ?? result.error ?? t('unknownError') })) + toast.error(t('importFailed', { name, error: reason ?? result.error ?? t('unknownError') }), { + id: `import-failed:${name}` + }) }, [t] ) diff --git a/src/renderer/src/components/notebook/note/NoteEditor.tsx b/src/renderer/src/components/notebook/note/NoteEditor.tsx index fb15b46..fe9f1a8 100644 --- a/src/renderer/src/components/notebook/note/NoteEditor.tsx +++ b/src/renderer/src/components/notebook/note/NoteEditor.tsx @@ -5,7 +5,6 @@ import StarterKit from '@tiptap/starter-kit' import { Markdown } from 'tiptap-markdown' import { Placeholder } from '@tiptap/extensions' import { ScrollArea } from '../../ui/scroll-area' -import { Toaster } from '../../ui/sonner' import { toast } from 'sonner' import { sourceDocumentExists } from '../../../../../shared/utils/citations' import { @@ -142,7 +141,7 @@ export default function NoteEditor({ noteId, content, onChange }: NoteEditorProp // 来源已删除：不导航，也不把锚点从正文里删掉。但要说明为什么没反应， // 否则「已禁用」和「坏了」在用户看来是一样的。 if (!sourceDocumentExists(context.documents, context.documentsLoaded, anchor.documentId)) { - toast.error(t('excerptSourceMissing')) + toast.error(t('excerptSourceMissing'), { id: 'excerpt-source-missing' }) return } context.openSourceAnchor(anchor) @@ -168,7 +167,6 @@ export default function NoteEditor({ noteId, content, onChange }: NoteEditorProp return (
- {/* 编辑器内容 */}
`, at the app root. + * + * sonner replays every still-active toast to a late subscriber, so a toaster that + * only exists while a note is open holds toasts fired from anywhere else (imports, + * saves, excerpts) and releases them in a burst when a note is next opened. This + * guard keeps a second mount from reappearing, wherever it is added. + */ + +const RENDERER_DIR = join('src', 'renderer', 'src') + +function sourceFiles(dir: string, files: string[] = []): string[] { + for (const entry of readdirSync(dir)) { + const path = join(dir, entry) + if (statSync(path).isDirectory()) sourceFiles(path, files) + else if (/\.(tsx|ts)$/.test(path)) files.push(path) + } + return files +} + +test('the app renders exactly one Toaster, at the root', () => { + const mounts: string[] = [] + + for (const file of sourceFiles(RENDERER_DIR)) { + const source = readFileSync(file, 'utf8') + if (source.includes(' was mounted somewhere other than the app root' + ) +}) + +test('the note editor does not own a toaster', () => { + const editor = readFileSync( + join(RENDERER_DIR, 'components', 'notebook', 'note', 'NoteEditor.tsx'), + 'utf8' + ) + assert.ok(!editor.includes(' { + const notePanel = readFileSync( + join(RENDERER_DIR, 'components', 'notebook', 'NotePanel.tsx'), + 'utf8' + ) + assert.ok( + notePanel.includes('id: `note-saved:${currentNote.id}`'), + 'note save must key its toast to the note' + ) + + const sourcePanel = readFileSync( + join(RENDERER_DIR, 'components', 'notebook', 'SourcePanel.tsx'), + 'utf8' + ) + assert.ok( + sourcePanel.includes('id: `excerpt:${currentNote.id}`'), + 'excerpt append must key its toast to the note' + ) + assert.ok( + sourcePanel.includes('id: `import-failed:${name}`'), + 'import failure must key its toast to the file' + ) + + const editor = readFileSync( + join(RENDERER_DIR, 'components', 'notebook', 'note', 'NoteEditor.tsx'), + 'utf8' + ) + assert.ok( + editor.includes("id: 'excerpt-source-missing'"), + 'the missing-source toast must be keyed' + ) +}) From 4546c106c0580170f8026aa556b444e35d3d4456 Mon Sep 17 00:00:00 2001 From: mrsibe Date: Tue, 29 Sep 2026 17:25:09 +0800 Subject: [PATCH 4/4] feat(chat): make a truncated answer recoverable in place (#179) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A reasoning model spends the output budget on thinking first, so the visible answer can end with finishReason 'length'. The notice was correct; there was no way out of it. - add the 截断时自动继续 setting and a continuation bound (default 2); on a truncated attempt the manager continues in place, seeding the next call with the accumulated answer, and stops at the bound - apply the same bound to the manual 继续生成 chain, recorded in the message metadata so it survives a reload - show the truncation notice and the recovery actions as one unit under the answer - extend the model connection with contextWindow, reasoning, reasoningEffort and reasoningBudget, and translate them per protocol; a connection that declares none sends exactly the request it did before --- src/main/config/defaults.ts | 4 + src/main/config/settingsManager.ts | 6 + src/main/ipc/chatHandlers.ts | 61 ++++++- src/main/ipc/validation.ts | 10 +- src/main/models/ModelClient.ts | 14 ++ .../models/protocols/anthropicMessages.ts | 21 +++ .../models/protocols/googleGenerativeAi.ts | 12 ++ .../models/protocols/openaiCompletions.ts | 11 ++ src/main/models/protocols/openaiResponses.ts | 7 + src/main/models/protocols/types.ts | 9 ++ src/main/services/chat/ChatStreamManager.ts | 108 ++++++++++++- src/preload/index.d.ts | 10 +- .../components/notebook/chat/MessageItem.tsx | 33 ++-- .../components/settings/GeneralSettings.tsx | 36 +++++ .../settings/ModelConnectionForm.tsx | 84 +++++++++- src/renderer/src/locales/en-US/settings.json | 16 ++ src/renderer/src/locales/en-US/ui.json | 3 +- src/renderer/src/locales/zh-CN/settings.json | 16 ++ src/renderer/src/locales/zh-CN/ui.json | 3 +- src/renderer/src/store/chatStore.ts | 13 +- src/shared/types/connection.ts | 33 ++++ src/shared/types/index.ts | 10 ++ test/chatStreamManager.test.ts | 149 ++++++++++++++++++ test/connectionCapabilities.test.ts | 137 ++++++++++++++++ 24 files changed, 769 insertions(+), 37 deletions(-) create mode 100644 test/connectionCapabilities.test.ts diff --git a/src/main/config/defaults.ts b/src/main/config/defaults.ts index dd9b878..33d695c 100644 --- a/src/main/config/defaults.ts +++ b/src/main/config/defaults.ts @@ -10,6 +10,10 @@ export const defaultSettings: AppSettings = { language: 'en-US', autoLaunch: false, hasCompletedOnboarding: false, + // Off by default: it spends provider tokens, so it is opt-in. Two continuations + // is the default bound when it is on (#179). + autoContinueOnTruncation: false, + maxAutoContinueAttempts: 2, prompts: { mindMap: { 'zh-CN': `你是知识结构分析专家,负责从笔记本内容中提炼核心知识结构。 diff --git a/src/main/config/settingsManager.ts b/src/main/config/settingsManager.ts index c68d106..6c8cc86 100644 --- a/src/main/config/settingsManager.ts +++ b/src/main/config/settingsManager.ts @@ -12,6 +12,12 @@ function mergeSettings(stored: Partial): AppSettings { language: stored.language ?? defaultSettings.language, autoLaunch: stored.autoLaunch ?? defaultSettings.autoLaunch, hasCompletedOnboarding: stored.hasCompletedOnboarding ?? defaultSettings.hasCompletedOnboarding, + // An older store has neither key; the defaults are what keeps them meaningful + // rather than `undefined` reaching the chat handler (#179). + autoContinueOnTruncation: + stored.autoContinueOnTruncation ?? defaultSettings.autoContinueOnTruncation, + maxAutoContinueAttempts: + stored.maxAutoContinueAttempts ?? defaultSettings.maxAutoContinueAttempts, prompts: { mindMap: { 'zh-CN': stored.prompts?.mindMap?.['zh-CN'] ?? defaultSettings.prompts!.mindMap!['zh-CN'], diff --git a/src/main/ipc/chatHandlers.ts b/src/main/ipc/chatHandlers.ts index 50abf80..1dcd81f 100644 --- a/src/main/ipc/chatHandlers.ts +++ b/src/main/ipc/chatHandlers.ts @@ -4,11 +4,12 @@ import { ConnectionManager } from '../models/ConnectionManager' import type { SessionAutoSwitchService } from '../services/SessionAutoSwitchService' import { KnowledgeService, toSearchResults } from '../services/KnowledgeService' import { buildRAGContext } from '../services/citations' -import { ChatStreamManager } from '../services/chat/ChatStreamManager' +import { ChatStreamManager, CONTINUE_INSTRUCTION } from '../services/chat/ChatStreamManager' import type { ChatTurnEvent } from '../../shared/types/chat' import { queriesTurnStore } from '../services/chat/turnStore' import { validateAndCleanMessages } from '../utils/messageValidator' import Logger from '../../shared/utils/logger' +import { settingsManager } from '../config' import type { AnswerSource, RetrievalSnapshot, RetrievalStatus } from '../../shared/types/chat' import type { Citation, CitationContext } from '../../shared/types/citation' import { parseRetrievalScope, scopeDocumentIds } from '../../shared/types/scope' @@ -17,17 +18,36 @@ import { ChatSchemas, validate } from './validation' /** * What a continuation is asked for. * - * English like the other internal prompts: it is an instruction to the model, not copy - * the reader sees. + * The constant itself is `ChatStreamManager`'s: automatic continuation (#179) and + * this manual path must ask for exactly the same thing. */ -const CONTINUE_INSTRUCTION = - 'Continue the answer from where it stopped. Do not repeat what has already been written, and do not restart it.' +const CONTINUE_MESSAGE = CONTINUE_INSTRUCTION /** * The renderer's view of a running turn: one channel, one shape. */ const TURN_EVENT_CHANNEL = 'chat:turn-event' +/** + * How many continuations one answer may spend, and whether a truncated one is + * continued automatically (#179). + * + * The manual「继续生成」chain and automatic continuation share the same bound, so a + * setting of 2 means "two more pieces of this answer", however they were asked for. + */ +function continuationSettings(): { limit: number; auto: boolean } { + const configured = settingsManager.getSettingSync('maxAutoContinueAttempts') + const limit = typeof configured === 'number' && configured > 0 ? Math.floor(configured) : 0 + const auto = settingsManager.getSettingSync('autoContinueOnTruncation') === true + return { limit, auto } +} + +/** How many continuations the persisted metadata says this answer already spent. */ +function continuationsSpent(metadata: unknown): number { + const value = (metadata as { continuations?: unknown } | null | undefined)?.continuations + return typeof value === 'number' && value > 0 ? Math.floor(value) : 0 +} + /** * Register chat-related IPC Handlers */ @@ -246,6 +266,7 @@ export function registerChatHandlers( // 4. 这一轮交给 ChatStreamManager（#140）。生命周期、唯一终态、落库与通知都在那里； // handler 只负责把 prompt 和检索结果准备好。 const client = await connectionManager.getChatClient() + const continuations = continuationSettings() const execution = streamManager.start({ sessionId, notebookId: session?.notebookId, @@ -257,6 +278,10 @@ export function registerChatHandlers( citations: retrieved.citations, citationContexts: retrieved.citationContexts, retrievalSnapshot: retrieved.retrievalSnapshot, + continuations: 0, + ...(continuations.auto && continuations.limit > 0 + ? { autoContinue: { limit: continuations.limit } } + : {}), emit: (turnEvent: ChatTurnEvent) => event.sender.send(TURN_EVENT_CHANNEL, turnEvent) }) @@ -300,6 +325,7 @@ export function registerChatHandlers( if (!prompt) return { success: false, error: 'Nothing to retry' } const client = await connectionManager.getChatClient() + const continuations = continuationSettings() const execution = streamManager.start({ sessionId: target.sessionId, notebookId: prompt.session?.notebookId, @@ -312,6 +338,11 @@ export function registerChatHandlers( citations: prompt.retrieved.citations, citationContexts: prompt.retrieved.citationContexts, retrievalSnapshot: prompt.retrieved.retrievalSnapshot, + // A sibling is a fresh answer, so its continuation budget restarts (#179). + continuations: 0, + ...(continuations.auto && continuations.limit > 0 + ? { autoContinue: { limit: continuations.limit } } + : {}), emit: (turnEvent: ChatTurnEvent) => event.sender.send(TURN_EVENT_CHANNEL, turnEvent) }) @@ -334,6 +365,18 @@ export function registerChatHandlers( const target = queries.getMessageById(messageId) if (!target) return { success: false, error: 'Message not found' } + const continuations = continuationSettings() + const spent = continuationsSpent(target.metadata) + // The bound is checked here as well as in the manager (#179): a manual chain + // must not be able to loop past the setting, however many times it is asked. + if (continuations.limit <= 0 || spent >= continuations.limit) { + return { + success: false, + reason: 'continuation-limit' as const, + error: `This answer has reached its continuation limit (${continuations.limit})` + } + } + const prompt = await replayPrompt(target) if (!prompt) return { success: false, error: 'Nothing to continue' } @@ -347,11 +390,17 @@ export function registerChatHandlers( reasoning: target.reasoningContent ?? '' }, userContent: prompt.question, + // This continuation is being spent now, so the count this turn records is the + // one already spent plus this one (#179). + continuations: spent + 1, + ...(continuations.auto && continuations.limit > 0 + ? { autoContinue: { limit: continuations.limit } } + : {}), messages: [ ...prompt.messages, // What had been written, then the instruction to carry on from it. { role: 'assistant', content: target.content }, - { role: 'user', content: CONTINUE_INSTRUCTION } + { role: 'user', content: CONTINUE_MESSAGE } ], client, retrieval: prompt.retrieved.retrieval, diff --git a/src/main/ipc/validation.ts b/src/main/ipc/validation.ts index e409938..11f802b 100644 --- a/src/main/ipc/validation.ts +++ b/src/main/ipc/validation.ts @@ -4,7 +4,7 @@ */ import { z } from 'zod' -import { API_PROTOCOLS, MODEL_CAPABILITIES } from '../../shared/types' +import { API_PROTOCOLS, MODEL_CAPABILITIES, REASONING_EFFORTS } from '../../shared/types' import { Result, Err, Ok } from '../../shared/types/result' import Logger from '../../shared/utils/logger' @@ -95,7 +95,13 @@ export const ConnectionSchemas = { apiKey: z.string(), modelId: z.string().min(1, 'Model ID 不能为空'), // Absent means "no ceiling": the provider's own default applies (#150). - maxOutputTokens: z.number().int().positive().optional() + maxOutputTokens: z.number().int().positive().optional(), + // Capability metadata (#179). All optional: a connection without any of it + // behaves exactly as connections did before they existed. + contextWindow: z.number().int().positive().optional(), + reasoning: z.boolean().optional(), + reasoningEffort: z.enum(REASONING_EFFORTS).optional(), + reasoningBudget: z.number().int().positive().optional() }) }), diff --git a/src/main/models/ModelClient.ts b/src/main/models/ModelClient.ts index b876ba3..975b663 100644 --- a/src/main/models/ModelClient.ts +++ b/src/main/models/ModelClient.ts @@ -7,6 +7,7 @@ import { embed, embedMany, streamText } from 'ai' import type { AsyncIterableStream, LanguageModel, LanguageModelUsage, UIMessageChunk } from 'ai' +import type { SharedV2ProviderOptions } from '@ai-sdk/provider' import type { APIMessage, ChatTokenUsage } from '../../shared/types/chat' import type { ModelCapability, ModelConnection } from '../../shared/types/connection' import { getProtocolAdapter } from './protocols' @@ -81,6 +82,14 @@ export class ModelClient { }) } + /** + * The provider options this connection's reasoning configuration translates to + * (#179), or undefined when nothing was declared. + */ + private chatProviderOptions(): SharedV2ProviderOptions | undefined { + return this.adapter.chatProviderOptions?.(this.connection) + } + /** * 流式发送消息 * @@ -93,6 +102,7 @@ export class ModelClient { Logger.debug('ModelClient', `Streaming with model: ${this.connection.modelId}`) const maxOutputTokens = this.connection.maxOutputTokens + const providerOptions = this.chatProviderOptions() const result = streamText({ model: this.getAIModel(), @@ -102,6 +112,10 @@ export class ModelClient { // default, which is what a reasoning model needs: its thinking shares this // budget with the answer, so 2048 of them was a ceiling reached early (#150). ...(maxOutputTokens !== undefined && { maxOutputTokens }), + // Only when the connection declares reasoning effort or a budget (#179). + // Absent keeps the request byte-for-byte what it was for an unknown + // OpenAI-compatible endpoint. + ...(providerOptions && { providerOptions }), // Owned by the caller, so that stopping a turn actually cancels the provider // request rather than only recording that it was stopped. abortSignal: options.signal diff --git a/src/main/models/protocols/anthropicMessages.ts b/src/main/models/protocols/anthropicMessages.ts index 7a9c611..be3bea0 100644 --- a/src/main/models/protocols/anthropicMessages.ts +++ b/src/main/models/protocols/anthropicMessages.ts @@ -5,11 +5,14 @@ */ import { createAnthropic } from '@ai-sdk/anthropic' +import type { SharedV2ProviderOptions } from '@ai-sdk/provider' import type { EmbeddingModel, LanguageModel } from 'ai' import type { ModelConnection } from '../../../shared/types/connection' import type { ProtocolAdapter } from './types' import { listAnthropicModels } from './modelListing' +const PROVIDER_OPTIONS_KEY = 'anthropic' + function createProvider(connection: ModelConnection) { return createAnthropic({ baseURL: connection.baseUrl, @@ -28,6 +31,24 @@ export const anthropicMessagesAdapter: ProtocolAdapter = { throw new Error('The anthropic-messages protocol does not support embeddings') }, + /** + * Anthropic takes an effort level, or an explicit thinking budget (#179). The + * budget form is used when the user gave one, because it is the more specific + * statement; otherwise the effort level is passed through. + */ + chatProviderOptions(connection): SharedV2ProviderOptions | undefined { + if (connection.reasoningBudget) { + return { + [PROVIDER_OPTIONS_KEY]: { + thinking: { type: 'enabled', budgetTokens: connection.reasoningBudget } + } + } + } + return connection.reasoningEffort + ? { [PROVIDER_OPTIONS_KEY]: { effort: connection.reasoningEffort } } + : undefined + }, + supportsModelListing: true, listModels: listAnthropicModels diff --git a/src/main/models/protocols/googleGenerativeAi.ts b/src/main/models/protocols/googleGenerativeAi.ts index 13ac38c..ba7b54e 100644 --- a/src/main/models/protocols/googleGenerativeAi.ts +++ b/src/main/models/protocols/googleGenerativeAi.ts @@ -34,6 +34,18 @@ export const googleGenerativeAiAdapter: ProtocolAdapter = { return dimensions ? { [PROVIDER_OPTIONS_KEY]: { outputDimensionality: dimensions } } : undefined }, + /** + * Gemini takes a thinking token budget rather than an effort level (#179), so + * only an explicit budget is sent. No budget means the provider's default. + */ + chatProviderOptions(connection) { + return connection.reasoningBudget + ? { + [PROVIDER_OPTIONS_KEY]: { thinkingConfig: { thinkingBudget: connection.reasoningBudget } } + } + : undefined + }, + supportsModelListing: true, listModels: listGoogleModels diff --git a/src/main/models/protocols/openaiCompletions.ts b/src/main/models/protocols/openaiCompletions.ts index 04be83b..8e14ba0 100644 --- a/src/main/models/protocols/openaiCompletions.ts +++ b/src/main/models/protocols/openaiCompletions.ts @@ -36,6 +36,17 @@ export const openaiCompletionsAdapter: ProtocolAdapter = { return dimensions ? { [PROVIDER_OPTIONS_KEY]: { dimensions } } : undefined }, + /** + * OpenAI-compatible reasoning effort (#179). The provider maps `low|medium|high` + * onto its own thinking budget; a model that does not reason ignores an unknown + * field, but we still only send it when the user asked for one. + */ + chatProviderOptions(connection) { + return connection.reasoningEffort + ? { [PROVIDER_OPTIONS_KEY]: { reasoningEffort: connection.reasoningEffort } } + : undefined + }, + supportsModelListing: true, listModels: listOpenAICompatibleModels diff --git a/src/main/models/protocols/openaiResponses.ts b/src/main/models/protocols/openaiResponses.ts index 85eb9b8..e4233ee 100644 --- a/src/main/models/protocols/openaiResponses.ts +++ b/src/main/models/protocols/openaiResponses.ts @@ -34,6 +34,13 @@ export const openaiResponsesAdapter: ProtocolAdapter = { return dimensions ? { [PROVIDER_OPTIONS_KEY]: { dimensions } } : undefined }, + /** OpenAI Responses takes the same effort vocabulary (#179). */ + chatProviderOptions(connection) { + return connection.reasoningEffort + ? { [PROVIDER_OPTIONS_KEY]: { reasoningEffort: connection.reasoningEffort } } + : undefined + }, + supportsModelListing: true, listModels: listOpenAICompatibleModels diff --git a/src/main/models/protocols/types.ts b/src/main/models/protocols/types.ts index bbfc976..659b548 100644 --- a/src/main/models/protocols/types.ts +++ b/src/main/models/protocols/types.ts @@ -31,6 +31,15 @@ export interface ProtocolAdapter { dimensions?: number ): SharedV2ProviderOptions | undefined + /** + * 把连接的推理配置翻译成该协议的 providerOptions（#179）。 + * + * 只在用户显式声明了 reasoningEffort / reasoningBudget 时返回东西；其余情况返回 + * undefined，请求体与引入这个能力之前完全一致 —— 未知的 OpenAI 兼容端点不会被 + * 塞进它不认识的字段。 + */ + chatProviderOptions?(connection: ModelConnection): SharedV2ProviderOptions | undefined + /** * 该协议是否支持从端点拉取模型列表（/v1/models 等）。 */ diff --git a/src/main/services/chat/ChatStreamManager.ts b/src/main/services/chat/ChatStreamManager.ts index 2ec42f6..4034003 100644 --- a/src/main/services/chat/ChatStreamManager.ts +++ b/src/main/services/chat/ChatStreamManager.ts @@ -77,6 +77,16 @@ const IDLE_TIMEOUT_MS = 60_000 /** Waits before the next attempt. */ const wait = (ms: number): Promise => new Promise((resolve) => setTimeout(resolve, ms)) +/** + * What an automatic continuation is asked for (#179). + * + * English like the manual continuation's instruction: it is a directive to the + * model, not copy the reader sees. One constant, because automatic and manual + * continuations must ask for exactly the same thing. + */ +export const CONTINUE_INSTRUCTION = + 'Continue the answer from where it stopped. Do not repeat what has already been written, and do not restart it.' + /** What the turn's persistence needs, as a port: see `turnStore.ts`. */ export interface ChatTurnStore { createTurn(sessionId: string, attemptOf?: string): { id: string } @@ -105,6 +115,18 @@ export interface ChatTurnRequest { resume?: { messageId: string; text: string; reasoning: string } /** The message the user sent, for the token estimate when no usage is reported. */ userContent: string + /** + * Automatic continuation after the output ceiling ends an answer (#179). + * + * Present only when the setting is on. `limit` is the total number of + * continuations this answer may spend, automatic and manual together. + */ + autoContinue?: { limit: number } + /** + * Continuations already spent on this answer, read from the persisted metadata + * (#179). A manual chain is bounded by the same counter. + */ + continuations?: number /** The prompt as it was assembled, retrieval included. Empty means nothing to send. */ messages: APIMessage[] /** Null when no chat model is configured. */ @@ -142,6 +164,8 @@ interface TurnEntry { timedOut?: boolean /** The attempt in flight, 1-based; logged, and what the retry bound counts. */ attempts?: number + /** How many continuations this answer has spent (#179). */ + continuations: number /** The watchdog for the attempt in flight. */ idle?: { reset: () => void; stop: () => void } /** Monotonic per execution, so a consumer can detect a gap. */ @@ -208,6 +232,9 @@ export class ChatStreamManager { sources: request.sources, citations: request.citations } + // How much of this answer's continuation budget is already spent (#179). Written + // with the provenance so a manual continue after a reload is still bounded. + if (request.continuations !== undefined) baseMetadata.continuations = request.continuations // The retrieval parameters are part of what an answer was built from, so they // are written with the sources, before the model is called (#157). if (request.retrievalSnapshot) baseMetadata.retrievalSnapshot = request.retrievalSnapshot @@ -222,7 +249,13 @@ export class ChatStreamManager { notebookId: request.notebookId }) - const entry: TurnEntry = { execution, request, baseMetadata, seq: 0 } + const entry: TurnEntry = { + execution, + request, + baseMetadata, + seq: 0, + continuations: request.continuations ?? 0 + } this.turns.set(execution.id, entry) this.turnIdByMessageId.set(execution.messageId, execution.id) @@ -275,20 +308,40 @@ export class ChatStreamManager { return executionId === undefined ? undefined : this.turns.get(executionId) } + private async run(entry: TurnEntry): Promise { + for (;;) { + const outcome = await this.runAttempts(entry) + + // A stop the reader asked for has already settled the turn. + if (entry.execution.isSettled) return + + // The output ceiling ended it and the answer may ask for more (#179): + // continue in place rather than settling a half answer. + if (outcome.status === 'truncated' && this.mayAutoContinue(entry)) { + this.beginContinuation(entry) + continue + } + + this.settleTurn(entry, outcome) + return + } + } + /** - * Run the turn, attempting it again while the retry policy says so (#150). + * Attempt the turn, retrying while the retry policy says so (#150). * * Each attempt is a fresh provider call with its own signal. A retry is only * reachable while nothing has reached the reader, so the answer on screen is never - * duplicated or silently replaced by a second attempt. + * duplicated or silently replaced by a second attempt. Returns the outcome the + * policy settled on; the caller decides whether that is the turn's end. */ - private async run(entry: TurnEntry): Promise { + private async runAttempts(entry: TurnEntry): Promise { for (let attempt = 1; ; attempt += 1) { entry.attempts = attempt const outcome = await this.attempt(entry) // A stop the reader asked for has already settled the turn. - if (entry.execution.isSettled) return + if (entry.execution.isSettled) return outcome if (shouldRetry({ outcome, attempt, contentSent: entry.contentSent === true })) { const delay = retryDelayMs(attempt + 1) @@ -299,15 +352,54 @@ export class ChatStreamManager { }); retrying in ${delay}ms` ) await wait(delay) - if (entry.execution.isSettled) return + if (entry.execution.isSettled) return outcome continue } - this.settleTurn(entry, outcome) - return + return outcome } } + /** Whether a truncated answer may be continued automatically (#179). */ + private mayAutoContinue(entry: TurnEntry): boolean { + const limit = entry.request.autoContinue?.limit ?? 0 + if (limit <= 0 || entry.continuations >= limit) return false + // Nothing to append to is not a continuation, and a turn that never reached the + // reader must not be continued silently. + return entry.contentSent === true && messageText(entry.message).length > 0 + } + + /** + * Seed the next provider call from what has arrived, plus the instruction to + * carry on (#179). + * + * The persisted answer keeps growing in the same message: the accumulated text is + * both the assembler's seed and the assistant turn handed back to the model, so a + * continuation cannot restart or duplicate the answer. + */ + private beginContinuation(entry: TurnEntry): void { + entry.continuations += 1 + const text = messageText(entry.message) + const reasoning = messageReasoning(entry.message) + + entry.request.resume = { messageId: entry.execution.messageId, text, reasoning } + entry.request.messages = [ + ...entry.request.messages, + { role: 'assistant', content: text }, + { role: 'user', content: CONTINUE_INSTRUCTION } + ] + entry.baseMetadata.continuations = entry.continuations + // The new call starts from nothing on the wire; only the seed it is given has + // content. Resetting lets the retry policy cover a continuation that fails + // before producing anything, which is the whole point of retrying it. + entry.contentSent = false + + Logger.info( + 'ChatStreamManager', + `Turn ${entry.execution.id} hit the output ceiling; continuing (${entry.continuations}/${entry.request.autoContinue?.limit ?? 0})` + ) + } + /** * One provider call, from opening the attempt to its outcome. * diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts index 9560239..66bb0bb 100644 --- a/src/preload/index.d.ts +++ b/src/preload/index.d.ts @@ -211,9 +211,13 @@ declare global { messageId: string ) => Promise<{ success: boolean; messageId?: string; error?: string }> /** Continue an answer that stopped, into the same message (#151). */ - continueMessage: ( - messageId: string - ) => Promise<{ success: boolean; messageId?: string; error?: string }> + continueMessage: (messageId: string) => Promise<{ + success: boolean + messageId?: string + error?: string + /** Why it was refused, when the reason is one the UI can translate (#179). */ + reason?: 'continuation-limit' + }> // 流式回合监听：SDK 自己的事件，加上结束这一轮的唯一 outcome（#141） onTurnEvent: (callback: (event: ChatTurnEvent) => void) => () => void diff --git a/src/renderer/src/components/notebook/chat/MessageItem.tsx b/src/renderer/src/components/notebook/chat/MessageItem.tsx index 2006d8e..e8a7c54 100644 --- a/src/renderer/src/components/notebook/chat/MessageItem.tsx +++ b/src/renderer/src/components/notebook/chat/MessageItem.tsx @@ -6,6 +6,7 @@ import rehypeHighlight from 'rehype-highlight' import rehypeKatex from 'rehype-katex' import { BookPlus, RotateCcw, Play } from 'lucide-react' import { useTranslation } from 'react-i18next' +import { toast } from 'sonner' import type { ChatMessage } from '../../../types/notebook' import ReasoningContent from './ReasoningContent' import AnswerSources from './AnswerSources' @@ -103,7 +104,14 @@ function MessageItem({ message }: MessageItemProps): ReactElement { // ask the same question again beside it. Both need the notebook the turn belongs to. const recoveryNotebookId = message.notebookId ?? currentNotebook?.id const handleContinue = (): void => { - if (recoveryNotebookId) void continueMessage(recoveryNotebookId, message.id) + if (!recoveryNotebookId) return + void continueMessage(recoveryNotebookId, message.id).then((result) => { + // The main process refuses past the continuation bound (#179); say why the + // button did nothing instead of leaving it silent. + if (!result.started && result.reason === 'continuation-limit') { + toast.error(t('ui:continueLimitReached')) + } + }) } const handleRetry = (): void => { if (recoveryNotebookId) void retryMessage(recoveryNotebookId, message.id) @@ -296,18 +304,14 @@ function MessageItem({ message }: MessageItemProps): ReactElement { {/* Citation coverage (#156). Stated only when something is not backed by a source; a normal answer does not need an "all good" line. */} {coverage && !isLive && } - {/* Why the answer stops where it does, when the model did not decide to stop - there itself. Beside the answer, never instead of it: the part that - arrived is what the reader was left with. */} + {/* Why the answer stops where it does, and the two ways out of it (#151, + #179), as one unit under the answer. Beside the answer, never instead of + it: the part that arrived is what the reader was left with. */} {!isLive && noticeKey && ( -

- {t(`chat:${noticeKey}`, { error: message.error?.message ?? '' })} -

- )} - {/* Action buttons - only shown when reply is complete and has content */} - {content && !isLive && ( -
- {/* Recovery, when the answer stopped early and asking again may help */} +
+

+ {t(`chat:${noticeKey}`, { error: message.error?.message ?? '' })} +

{canRecover(message.status) && ( <>
+ )} + {/* Action buttons - only shown when reply is complete and has content */} + {content && !isLive && ( +
{/* Copy button */}
)} + {/* Capability metadata (#179). Every field is optional, and a connection + without any of them sends exactly the same request as before. */} + {capability === 'chat' && ( +
+ +

{t('reasoningHint')}

+ +
+ + +

{t('reasoningEffortHint')}

+
+ +
+ + update({ reasoningBudget: toPositive(e.target.value) })} + /> +

{t('reasoningBudgetHint')}

+
+ +
+ + update({ contextWindow: toPositive(e.target.value) })} + /> +

{t('contextWindowHint')}

+
+
+ )} + {/* Test connection */}