feat(import): add safe batch coordinator

This commit is contained in:
digua committed 2026-07-31 20:22:50 +08:00
1 parent 4155386676
commit 62b7104f1f
31 files changed
+2455 -223

No files matched your search

+133
View File
@@ -1,4 +1,5 @@
import assert from 'node:assert/strict'
import fs from 'node:fs'
import { describe, it } from 'node:test'
import Fastify from 'fastify'
import multipart from '@fastify/multipart'
@@ -60,7 +61,139 @@ function directoryMultipartPayload(): { payload: Buffer; contentType: string } {
}
}
function batchMultipartPayload(): { payload: Buffer; contentType: string } {
const boundary = '----chatlab-batch-import-route-test'
const field = [`--${boundary}`, 'Content-Disposition: form-data; name="sessionGapThreshold"', '', '7200'].join('\r\n')
const files = ['first.json', 'second.json'].map((filename, index) =>
[
`--${boundary}`,
`Content-Disposition: form-data; name="files"; filename="${filename}"`,
'Content-Type: application/json',
'',
JSON.stringify({ index }),
].join('\r\n')
)
return {
payload: Buffer.from([field, ...files, `--${boundary}--`, ''].join('\r\n')),
contentType: `multipart/form-data; boundary=${boundary}`,
}
}
describe('CLI Web automatic import route', () => {
it('runs uploaded files through one backend batch coordinator', async () => {
const app = Fastify()
const calls: Array<{ ids: string[]; paths: string[]; threshold?: number; concurrency?: number }> = []
try {
await app.register(multipart)
registerImportRoutes(
app,
{} as any,
{
runAutoImportBatch: async (items: any[], options: any) => {
calls.push({
ids: items.map((item) => item.id),
paths: items.map((item) => item.filePath),
threshold: options.sessionGapThreshold,
concurrency: options.concurrency,
})
options.onItemStart(items[0], 0)
options.onItemProgress(items[0], 0, { stage: 'parsing', percentage: 50, message: '' })
const results = [
{ id: '0', status: 'success', result: { success: true, sessionId: 'session-1' } },
{ id: '1', status: 'failed', error: 'expected failure' },
]
options.onItemComplete(items[0], 0, results[0])
options.onItemComplete(items[1], 1, results[1])
return results
},
} as any
)
const body = batchMultipartPayload()
const response = await app.inject({
method: 'POST',
url: '/_web/import/batch',
headers: { 'content-type': body.contentType },
payload: body.payload,
})
assert.equal(response.statusCode, 200)
assert.deepEqual(
calls.map(({ ids, threshold, concurrency }) => ({ ids, threshold, concurrency })),
[{ ids: ['0', '1'], threshold: 7200, concurrency: undefined }]
)
assert.equal(
calls[0].paths.every((filePath) => !fs.existsSync(filePath)),
true
)
assert.match(response.body, /event: batch-start/)
assert.match(response.body, /event: batch-progress/)
assert.match(response.body, /event: batch-complete/)
assert.match(response.body, /event: done/)
assert.match(response.body, /"sessionId":"session-1"/)
assert.match(response.body, /"error":"expected failure"/)
} finally {
await app.close()
}
})
it('cancels pending Web batch items without disconnecting the result stream', async () => {
const app = Fastify()
let notifyStarted: (() => void) | undefined
const started = new Promise<void>((resolve) => {
notifyStarted = resolve
})
try {
await app.register(multipart)
registerImportRoutes(
app,
{} as any,
{
runAutoImportBatch: async (items: any[], options: any) => {
notifyStarted?.()
await new Promise<void>((resolve) => {
if (options.signal.aborted) resolve()
else options.signal.addEventListener('abort', () => resolve(), { once: true })
})
const results = [
{ id: '0', status: 'success', result: { success: true, sessionId: 'session-1' } },
{ id: '1', status: 'cancelled' },
]
options.onItemComplete(items[0], 0, results[0])
options.onItemComplete(items[1], 1, results[1])
return results
},
} as any
)
const body = batchMultipartPayload()
const responsePromise = app.inject({
method: 'POST',
url: '/_web/import/batch',
headers: {
'content-type': body.contentType,
'x-chatlab-import-batch-id': 'batch-cancel-test',
},
payload: body.payload,
})
await started
const cancelResponse = await app.inject({
method: 'POST',
url: '/_web/import/batch/batch-cancel-test/cancel',
})
const response = await responsePromise
assert.equal(cancelResponse.statusCode, 200)
assert.deepEqual(cancelResponse.json(), { success: true, active: true })
assert.match(response.body, /"id":"0","status":"success"/)
assert.match(response.body, /"id":"1","status":"cancelled"/)
assert.match(response.body, /event: done/)
} finally {
await app.close()
}
})
it('forwards parser options and preserves a zero-new incremental result in the done event', async () => {
const app = Fastify()
const calls: Array<{ formatId?: string; chatIndex?: number; sessionGapThreshold?: number }> = []
+90 -1
View File
@@ -13,6 +13,7 @@ import {
} from '@openchatlab/node-runtime'
import {
autoImport,
autoImportBatch,
streamImport,
incrementalImport,
analyzeIncrementalImport,
@@ -23,7 +24,12 @@ import {
scanMultiChatFile,
findEntryFileInDirectory,
} from '../../../import'
import type { AutoImportResult, StreamImportOptions } from '../../../import'
import type {
AutoImportBatchOptions,
AutoImportBatchRequest,
AutoImportResult,
StreamImportOptions,
} from '../../../import'
import { resolveNativeBinding } from './helpers'
const ARCHIVE_UPLOAD_LIMIT = 50 * 1024 * 1024 * 1024
@@ -31,6 +37,10 @@ const ARCHIVE_UPLOAD_LIMIT = 50 * 1024 * 1024 * 1024
interface ImportRouteOptions {
sourceManager?: ArchiveImportSourceManager
runAutoImport?: (filePath: string, options: StreamImportOptions) => Promise<AutoImportResult>
runAutoImportBatch?: (
items: AutoImportBatchRequest[],
options: AutoImportBatchOptions
) => ReturnType<typeof autoImportBatch>
runPreparedImport?: (
manifestPath: string,
onProgress: (progress: unknown) => void,
@@ -65,8 +75,10 @@ export function registerImportRoutes(
dbManager: DatabaseManager,
options: ImportRouteOptions = {}
): void {
const activeBatchControllers = new Map<string, AbortController>()
const sourceManager = options.sourceManager ?? new ArchiveImportSourceManager()
const runAutoImport = options.runAutoImport ?? autoImport.bind(null, dbManager)
const runAutoImportBatch = options.runAutoImportBatch ?? autoImportBatch.bind(null, dbManager)
const runPreparedImport =
options.runPreparedImport ??
(async (manifestPath: string, onProgress: (progress: unknown) => void, sessionGapThreshold?: number) => {
@@ -252,6 +264,83 @@ export function registerImportRoutes(
}
})
server.post('/_web/import/batch', async (request, reply) => {
const parts = (request as any).parts()
if (!parts) return reply.code(400).send({ error: 'No files uploaded' })
const requestedBatchId = request.headers['x-chatlab-import-batch-id']
const batchId = typeof requestedBatchId === 'string' && requestedBatchId ? requestedBatchId : randomUUID()
if (activeBatchControllers.has(batchId)) {
return reply.code(409).send({ error: 'Batch import id is already active' })
}
const tmpDir = createChatLabTempDir('imports', 'batch-')
const items: AutoImportBatchRequest[] = []
let sessionGapThreshold: number | undefined
let completed = false
const controller = new AbortController()
activeBatchControllers.set(batchId, controller)
try {
for await (const part of parts) {
if (part.type === 'field' && part.fieldname === 'sessionGapThreshold') {
sessionGapThreshold = parseOptionalInteger(part.value)
continue
}
if (part.type !== 'file') continue
const index = items.length
const safeName = path.basename(part.filename || `upload-${index}`)
const filePath = path.join(tmpDir, `${index}-${safeName}`)
await pipeline(part.file, fs.createWriteStream(filePath, { flags: 'wx' }))
items.push({ id: String(index), filePath })
}
if (items.length === 0) return reply.code(400).send({ error: 'No files uploaded' })
reply.raw.writeHead(200, {
'Content-Type': 'text/event-stream',
'Cache-Control': 'no-cache',
Connection: 'keep-alive',
})
reply.raw.once('close', () => {
if (!completed) controller.abort()
})
const sendEvent = (event: string, eventData: unknown) => {
if (!reply.raw.destroyed) {
reply.raw.write(`event: ${event}\ndata: ${JSON.stringify(eventData)}\n\n`)
}
}
const results = await runAutoImportBatch(items, {
sessionGapThreshold,
signal: controller.signal,
onItemStart: (_item, index) => sendEvent('batch-start', { index }),
onItemProgress: (_item, index, progress) => sendEvent('batch-progress', { index, progress }),
onItemComplete: (_item, index, result) => sendEvent('batch-complete', { index, result }),
})
completed = true
sendEvent('done', results)
} catch (error) {
const payload = { error: error instanceof Error ? error.message : String(error) }
if (!reply.raw.headersSent) return reply.code(500).send(payload)
if (!reply.raw.destroyed) {
reply.raw.write(`event: error\ndata: ${JSON.stringify(payload)}\n\n`)
}
} finally {
completed = true
if (activeBatchControllers.get(batchId) === controller) activeBatchControllers.delete(batchId)
reply.raw.end()
cleanupTemp(tmpDir)
}
})
server.post<{ Params: { batchId: string } }>('/_web/import/batch/:batchId/cancel', async (request) => {
const controller = activeBatchControllers.get(request.params.batchId)
if (!controller) return { success: false, active: false }
controller.abort()
return { success: true, active: true }
})
// ==================== Directory Import ====================
server.post('/_web/import-directory', async (request, reply) => {
+3
View File
@@ -2,6 +2,7 @@
export {
streamImport,
autoImport,
autoImportBatch,
analyzeAutoImport,
incrementalImport,
analyzeIncrementalImport,
@@ -24,4 +25,6 @@ export type {
AnalyzeNewImportResult,
AutoImportAnalysisResult,
AutoImportResult,
AutoImportBatchRequest,
AutoImportBatchOptions,
} from './stream-import'
+122 -1
View File
@@ -12,7 +12,7 @@ import {
readDataDirCompatibilityMeta,
withDataDirImportLock,
} from '@openchatlab/node-runtime'
import { analyzeAutoImport, autoImport, streamImport } from './stream-import'
import { analyzeAutoImport, autoImport, autoImportBatch, streamImport } from './stream-import'
const nativeBinding = path.resolve('apps/cli/native/better_sqlite3.node')
@@ -67,6 +67,20 @@ function writeTempDuplicateChatFile(dir: string): string {
return filePath
}
function writeBatchChatFile(dir: string, filename: string, groupId: string, content: string): string {
const filePath = path.join(dir, filename)
fs.writeFileSync(
filePath,
JSON.stringify({
chatlab: { version: '0.0.2', exportedAt: 1711468800 },
meta: { name: groupId, platform: 'qq', type: 'group', groupId },
members: [{ platformId: 'u1', accountName: 'Alice' }],
messages: [{ sender: 'u1', accountName: 'Alice', timestamp: 1711468800, type: 0, content }],
})
)
return filePath
}
test('streamImport raises the data directory gate after creating a current-schema database', async () => {
const root = makeTempDir()
fs.mkdirSync(path.join(root, 'data', 'databases'), { recursive: true })
@@ -151,6 +165,113 @@ test('autoImport creates once and then incrementally imports the same stable cha
assert.equal(manager.listSessionIds().length, 1)
})
test('autoImportBatch holds one lock, coalesces the same target, and preserves independent results', async (t) => {
const root = makeTempDir()
t.after(() => fs.rmSync(root, { recursive: true, force: true }))
fs.mkdirSync(path.join(root, 'data', 'databases'), { recursive: true })
const manager = new DatabaseManager(createPathProvider(root), {
nativeBinding,
runtime: { version: '0.25.1', kind: 'cli' },
})
const first = writeBatchChatFile(root, 'group-a-1.json', 'group-a', 'first')
const second = writeBatchChatFile(root, 'group-a-2.json', 'group-a', 'second')
const independent = writeBatchChatFile(root, 'group-b.json', 'group-b', 'independent')
const results = await autoImportBatch(
manager,
[
{ id: 'a-1', filePath: first },
{ id: 'b', filePath: independent },
{ id: 'a-2', filePath: second },
],
{ concurrency: 2, sessionGapThreshold: 7200 }
)
assert.deepEqual(
results.map((result) => result.status),
['success', 'success', 'success']
)
assert.equal(results[0].status === 'success' && results[0].result.importMode, 'created')
assert.equal(results[2].status === 'success' && results[2].result.importMode, 'incremental')
assert.equal(
results[0].status === 'success' &&
results[2].status === 'success' &&
results[0].result.sessionId === results[2].result.sessionId,
true
)
assert.equal(manager.listSessionIds().length, 2)
const sameTargetId = results[0].status === 'success' ? results[0].result.sessionId! : ''
const db = manager.openRawSessionDatabase(sameTargetId, { readonly: true })
const row = db
.prepare(
`SELECT
(SELECT COUNT(*) FROM message) AS messages,
(SELECT session_gap_threshold FROM meta LIMIT 1) AS gapThreshold`
)
.get() as { messages: number; gapThreshold: number }
db.close()
assert.deepEqual(row, { messages: 2, gapThreshold: 7200 })
})
test('autoImportBatch rejects an external writer and cancels work that has not started', async (t) => {
const root = makeTempDir()
t.after(() => fs.rmSync(root, { recursive: true, force: true }))
fs.mkdirSync(path.join(root, 'data', 'databases'), { recursive: true })
const manager = new DatabaseManager(createPathProvider(root), {
nativeBinding,
runtime: { version: '0.25.1', kind: 'cli' },
})
const file = writeBatchChatFile(root, 'locked.json', 'locked', 'locked')
const lockedResults = await withDataDirImportLock(manager.getUserDataDir(), () =>
autoImportBatch(manager, [{ id: 'locked', filePath: file }])
)
assert.deepEqual(lockedResults, [{ id: 'locked', status: 'failed', error: IMPORT_IN_PROGRESS_ERROR_KEY }])
const controller = new AbortController()
const items = Array.from({ length: 5 }, (_, index) => ({
id: String(index),
filePath: writeBatchChatFile(root, `cancel-${index}.json`, `cancel-${index}`, `message-${index}`),
sessionId: `cancel-${index}`,
}))
let completed = 0
const cancelledResults = await autoImportBatch(manager, items, {
concurrency: 1,
signal: controller.signal,
onItemComplete() {
completed++
if (completed === 1) controller.abort()
},
})
assert.equal(cancelledResults[0].status, 'success')
assert.deepEqual(
cancelledResults.slice(1).map((result) => result.status),
['cancelled', 'cancelled', 'cancelled', 'cancelled']
)
assert.equal(manager.listSessionIds().length, 1)
})
test('autoImportBatch removes newly created sessions when the shared compatibility gate fails', async (t) => {
const root = makeTempDir()
t.after(() => fs.rmSync(root, { recursive: true, force: true }))
fs.mkdirSync(path.join(root, 'data', 'databases'), { recursive: true })
const manager = new DatabaseManager(createPathProvider(root), {
nativeBinding,
runtime: { version: '0.25.1', kind: 'cli' },
})
manager.raiseCurrentChatDbCompatibilityGate = () => {
throw new Error('compatibility gate unavailable')
}
const file = writeBatchChatFile(root, 'gate.json', 'gate', 'message')
const results = await autoImportBatch(manager, [{ id: 'gate', filePath: file, sessionId: 'gate-session' }])
assert.equal(results[0].status, 'failed')
assert.equal(results[0].status === 'failed' && results[0].error, 'compatibility gate unavailable')
assert.equal(fs.existsSync(path.join(root, 'data', 'databases', 'gate-session.db')), false)
})
test('analyzeAutoImport previews a new session without writing a database', async (t) => {
const root = makeTempDir()
t.after(() => fs.rmSync(root, { recursive: true, force: true }))
+194 -56
View File
@@ -11,7 +11,9 @@ import {
DataDirCompatibilityError,
IMPORT_IN_PROGRESS_ERROR_KEY,
ImportInProgressError,
resolveDefaultBatchConcurrency,
analyzeAutoImportFile as sharedAnalyzeAutoImportFile,
autoImportBatch as sharedAutoImportBatch,
autoImportFile as sharedAutoImportFile,
streamingImport,
analyzeNewImport as sharedAnalyzeNewImport,
@@ -29,6 +31,7 @@ import type {
ImportOptions,
AnalyzeNewImportResult,
AutoImportAnalysisResult,
AutoImportBatchItemResult,
AutoImportResult,
} from '@openchatlab/node-runtime'
import {
@@ -91,43 +94,45 @@ function buildStreamImportDeps(
function createProgressAdapter(onProgress?: (progress: StreamImportProgress) => void): ImportProgressCallback {
if (!onProgress) return () => {}
return (progress) => {
let stage: StreamImportProgress['stage'] = 'parsing'
let pct = 0
switch (progress.stage) {
case 'detecting':
stage = 'detecting'
pct = 5
break
case 'parsing':
stage = 'parsing'
pct = Math.min(Math.round(progress.percentage * 0.7), 70)
break
case 'saving':
stage = 'saving'
pct = 80
break
case 'indexing':
stage = 'indexing'
pct = 90
break
case 'done':
stage = 'done'
pct = 100
break
case 'error':
stage = 'error'
pct = 0
break
}
onProgress({
stage,
progress: pct,
message: progress.message || '',
bytesRead: progress.bytesRead,
totalBytes: progress.totalBytes,
messagesProcessed: progress.messagesProcessed,
})
return (progress) => onProgress(createProgressAdapterValue(progress))
}
function createProgressAdapterValue(progress: Parameters<ImportProgressCallback>[0]): StreamImportProgress {
let stage: StreamImportProgress['stage'] = 'parsing'
let pct = 0
switch (progress.stage) {
case 'detecting':
stage = 'detecting'
pct = 5
break
case 'parsing':
stage = 'parsing'
pct = Math.min(Math.round(progress.percentage * 0.7), 70)
break
case 'saving':
stage = 'saving'
pct = 80
break
case 'indexing':
stage = 'indexing'
pct = 90
break
case 'done':
stage = 'done'
pct = 100
break
case 'error':
stage = 'error'
pct = 0
break
}
return {
stage,
progress: pct,
message: progress.message || '',
bytesRead: progress.bytesRead,
totalBytes: progress.totalBytes,
messagesProcessed: progress.messagesProcessed,
}
}
@@ -142,7 +147,8 @@ function deleteSessionDatabase(dbManager: DatabaseManager, sessionId: string): v
async function streamImportUnlocked(
dbManager: DatabaseManager,
filePath: string,
options?: StreamImportOptions
options?: StreamImportOptions,
updateCompatibilityGate = true
): Promise<StreamImportResult> {
const { formatId, chatIndex, sessionGapThreshold, onProgress, sessionId } = options || {}
@@ -156,6 +162,8 @@ async function streamImportUnlocked(
const result = await streamingImport(filePath, deps, formatOptions, sessionId)
if (!result.success || !result.sessionId) return result
if (!updateCompatibilityGate) return result
try {
dbManager.raiseCurrentChatDbCompatibilityGate()
} catch (error) {
@@ -191,6 +199,24 @@ export async function autoImport(
dbManager: DatabaseManager,
filePath: string,
options?: StreamImportOptions
): Promise<AutoImportResult> {
try {
return await withDataDirImportLock(dbManager.getUserDataDir(), () =>
autoImportUnlocked(dbManager, filePath, options)
)
} catch (error) {
if (error instanceof ImportInProgressError) {
return { success: false, error: IMPORT_IN_PROGRESS_ERROR_KEY }
}
throw error
}
}
async function autoImportUnlocked(
dbManager: DatabaseManager,
filePath: string,
options?: StreamImportOptions,
updateCompatibilityGate = true
): Promise<AutoImportResult> {
const { formatId, chatIndex, onProgress, sessionId } = options ?? {}
const formatOptions: Record<string, unknown> = {}
@@ -198,37 +224,147 @@ export async function autoImport(
if (chatIndex !== undefined) formatOptions.chatIndex = chatIndex
const progressAdapter = createProgressAdapter(onProgress)
return sharedAutoImportFile(
filePath,
{
listSessionIds: () => dbManager.listSessionIds(),
openReadonly: (candidateSessionId) => dbManager.openRawSessionDatabase(candidateSessionId, { readonly: true }),
onProgress: progressAdapter,
sessionExists: (candidateSessionId) => dbManager.listSessionIds().includes(candidateSessionId),
createSession: (sourcePath, sourceFormatOptions, explicitSessionId) =>
streamImportUnlocked(
dbManager,
sourcePath,
{
...options,
...sourceFormatOptions,
sessionId: explicitSessionId,
},
updateCompatibilityGate
),
appendSession: (targetSessionId, sourcePath, sourceFormatOptions) =>
incrementalImportUnlocked(
dbManager,
targetSessionId,
sourcePath,
{
...sourceFormatOptions,
onProgress: progressAdapter,
},
updateCompatibilityGate
),
},
{
explicitSessionId: sessionId,
formatOptions,
}
)
}
export interface AutoImportBatchRequest {
id: string
filePath: string
formatId?: string
chatIndex?: number
sessionId?: string
}
export interface AutoImportBatchOptions {
concurrency?: number
sessionGapThreshold?: number
signal?: AbortSignal
onItemStart?: (item: AutoImportBatchRequest, index: number) => void
onItemProgress?: (item: AutoImportBatchRequest, index: number, progress: StreamImportProgress) => void
onItemComplete?: (item: AutoImportBatchRequest, index: number, result: AutoImportBatchItemResult) => void
}
export async function autoImportBatch(
dbManager: DatabaseManager,
items: AutoImportBatchRequest[],
options: AutoImportBatchOptions = {}
): Promise<AutoImportBatchItemResult[]> {
try {
return await withDataDirImportLock(dbManager.getUserDataDir(), () =>
sharedAutoImportFile(
filePath,
return await withDataDirImportLock(dbManager.getUserDataDir(), async () => {
const results = await sharedAutoImportBatch(
items.map((item) => ({
id: item.id,
filePath: item.filePath,
options: {
explicitSessionId: item.sessionId,
formatOptions: {
...(item.formatId ? { formatId: item.formatId } : {}),
...(item.chatIndex !== undefined ? { chatIndex: item.chatIndex } : {}),
},
},
})),
{
listSessionIds: () => dbManager.listSessionIds(),
openReadonly: (candidateSessionId) =>
dbManager.openRawSessionDatabase(candidateSessionId, { readonly: true }),
onProgress: progressAdapter,
sessionExists: (candidateSessionId) => dbManager.listSessionIds().includes(candidateSessionId),
createSession: (sourcePath, sourceFormatOptions, explicitSessionId) =>
streamImportUnlocked(dbManager, sourcePath, {
...options,
...sourceFormatOptions,
sessionId: explicitSessionId,
}),
streamImportUnlocked(
dbManager,
sourcePath,
{
...sourceFormatOptions,
sessionId: explicitSessionId,
sessionGapThreshold: options.sessionGapThreshold,
},
false
),
appendSession: (targetSessionId, sourcePath, sourceFormatOptions) =>
incrementalImportUnlocked(dbManager, targetSessionId, sourcePath, {
...sourceFormatOptions,
onProgress: progressAdapter,
}),
incrementalImportUnlocked(
dbManager,
targetSessionId,
sourcePath,
{
...sourceFormatOptions,
},
false
),
},
{
explicitSessionId: sessionId,
formatOptions,
concurrency: options.concurrency ?? resolveDefaultBatchConcurrency(items.length),
signal: options.signal,
onItemStart: (_item, index) => options.onItemStart?.(items[index], index),
onItemProgress: (_item, index, progress) =>
options.onItemProgress?.(items[index], index, createProgressAdapterValue(progress)),
onItemComplete: (_item, index, result) => options.onItemComplete?.(items[index], index, result),
}
)
)
if (results.some((result) => result.status === 'success')) {
try {
dbManager.raiseCurrentChatDbCompatibilityGate()
} catch (error) {
const message = error instanceof Error ? error.message : String(error)
for (const result of results) {
if (result.status === 'success' && result.result.importMode === 'created' && result.result.sessionId) {
deleteSessionDatabase(dbManager, result.result.sessionId)
}
}
return results.map((result) =>
result.status === 'success'
? {
id: result.id,
status: 'failed' as const,
error: message,
result: { ...result.result, success: false, error: message },
}
: result
)
}
}
return results
})
} catch (error) {
if (error instanceof ImportInProgressError) {
return { success: false, error: IMPORT_IN_PROGRESS_ERROR_KEY }
return items.map((item) => ({
id: item.id,
status: 'failed' as const,
error: IMPORT_IN_PROGRESS_ERROR_KEY,
}))
}
throw error
}
@@ -296,7 +432,8 @@ async function incrementalImportUnlocked(
dbManager: DatabaseManager,
sessionId: string,
filePath: string,
options?: ImportOptions & { onProgress?: ImportProgressCallback }
options?: ImportOptions & { onProgress?: ImportProgressCallback },
updateCompatibilityGate = true
): Promise<IncrementalImportResult> {
const { onProgress, ...importOpts } = options || {}
let compatibilityError: DataDirCompatibilityError | null = null
@@ -310,6 +447,7 @@ async function incrementalImportUnlocked(
)
if (compatibilityError) throw compatibilityError
if (!result.success) return result
if (!updateCompatibilityGate) return result
try {
dbManager.raiseCurrentChatDbCompatibilityGate()
+36
View File
@@ -233,6 +233,42 @@ export function registerChatHandlers(ctx: IpcContext): void {
}
})
ipcMain.handle(
'chat:importBatch',
async (
_,
batchId: string,
items: Array<{ id: string; filePath: string }>,
options?: { concurrency?: number; sessionGapThreshold?: number }
) => {
try {
const results = await worker.autoImportBatch(
batchId,
items,
(progress) => {
win.webContents.send('chat:importBatchProgress', { batchId, ...progress })
},
options
)
if (results.some((result) => result.status === 'success')) {
win.webContents.send('api:importCompleted')
}
return results
} catch (error) {
return items.map((item) => ({
id: item.id,
status: 'failed' as const,
error: error instanceof Error ? error.message : String(error),
}))
}
}
)
ipcMain.handle('chat:cancelImportBatch', async (_, batchId: string) => {
await worker.cancelAutoImportBatch(batchId)
return { success: true }
})
// ==================== 增量导入 ====================
ipcMain.handle('chat:analyzeIncrementalImport', async (_, sessionId: string, filePath: string) => {
+16
View File
@@ -81,6 +81,7 @@ import {
import {
streamImport,
autoImport,
autoImportBatch,
streamParseFileInfo,
analyzeIncrementalImport,
incrementalImport,
@@ -89,6 +90,8 @@ import {
pushImport,
} from './import'
const importBatchControllers = new Map<string, AbortController>()
initDbDir(workerData.dbDir, workerData.cacheDir, workerData.tempDir, workerData.nativeBinding, workerData.logsDir)
// 初始化 NLP 词库目录
@@ -191,6 +194,10 @@ const syncHandlers: Record<string, (payload: any) => any> = {
closeAllDatabases()
return true
},
cancelAutoImportBatch: (p) => {
importBatchControllers.get(p.batchId)?.abort()
return true
},
// 成员管理
getMembers: (p) => getMembers(p.sessionId),
@@ -253,6 +260,15 @@ const asyncHandlers: Record<string, (payload: any, requestId: string) => Promise
// 流式导入
streamImport: (p, id) => streamImport(p.filePath, id, p.formatOptions, p.externalSessionId, p.sessionGapThreshold),
autoImport: (p, id) => autoImport(p.filePath, id, p.formatOptions, p.explicitSessionId, p.sessionGapThreshold),
autoImportBatch: async (p, id) => {
const controller = new AbortController()
importBatchControllers.set(p.batchId, controller)
try {
return await autoImportBatch(p.items, id, p.concurrency, p.sessionGapThreshold, controller.signal)
} finally {
importBatchControllers.delete(p.batchId)
}
},
analyzePushImport: (p) => analyzePushImport(p.sessionId, p.payload),
pushImport: (p) => pushImport(p.sessionId, p.payload),
// 流式解析文件信息(用于合并预览)
+3
View File
@@ -7,10 +7,13 @@
export {
streamImport,
autoImport,
autoImportBatch,
streamParseFileInfo,
analyzeNewImport,
type StreamImportResult,
type AutoImportResult,
type AutoImportBatchItemResult,
type WorkerAutoImportBatchItem,
type StreamParseFileInfoResult,
type AnalyzeNewImportResult,
} from './streamImport'
+93 -19
View File
@@ -11,28 +11,31 @@ import * as path from 'path'
import type Database from 'better-sqlite3'
import {
BetterSqliteAdapter,
autoImportBatch as sharedAutoImportBatch,
autoImportFile as sharedAutoImportFile,
streamingImport,
analyzeNewImport as sharedAnalyzeNewImport,
streamParseFileInfo as sharedStreamParseFileInfo,
TEMP_DB_SCHEMA,
computeAndSetOverviewCache,
createImportPerfLogger,
deleteSessionCache,
getCurrentLogFile,
logError,
logInfo,
logPerf,
logPerfDetail,
logSummary,
resetPerfLog,
} from '@openchatlab/node-runtime'
import type { AutoImportResult, StreamImportDeps, StreamImportResult, ImportLogger } from '@openchatlab/node-runtime'
import type {
AutoImportBatchItemResult,
AutoImportResult,
StreamImportDeps,
StreamImportResult,
ImportLogger,
} from '@openchatlab/node-runtime'
import { sendProgress, generateSessionId, getDbPath, createDatabaseWithoutIndexes } from './utils'
import { incrementalImport } from './incrementalImport'
import { getCacheDir, getDbDir, getTempDir, openRawDatabase, initPerfLog } from '../core'
import { getCacheDir, getDbDir, getLogsDir, getTempDir, openRawDatabase } from '../core'
import { getImportLogDir } from '../core/perfLogPath'
export type { StreamImportResult }
export type { AutoImportResult }
export type { AutoImportBatchItemResult }
export type { AnalyzeNewImportResult, StreamParseFileInfoResult } from '@openchatlab/node-runtime'
export type { SkipReasons, ImportDiagnostics } from '@openchatlab/node-runtime'
@@ -54,16 +57,7 @@ function createTempDatabase(dbPath: string): Database.Database {
}
function buildElectronLogger(): ImportLogger {
return {
info: logInfo,
error: (message: string, err?: Error) => logError(message, err),
perf: (label: string, messageCount: number, batchSize?: number) => logPerf(label, messageCount, batchSize),
perfDetail: logPerfDetail,
summary: logSummary,
reset: resetPerfLog,
init: initPerfLog,
getCurrentLogFile,
}
return createImportPerfLogger(getImportLogDir(getLogsDir()))
}
function buildStreamImportDeps(requestId: string, sessionGapThreshold?: number): StreamImportDeps {
@@ -155,6 +149,86 @@ export async function autoImport(
)
}
export interface WorkerAutoImportBatchItem {
id: string
filePath: string
formatOptions?: Record<string, unknown>
explicitSessionId?: string
}
export async function autoImportBatch(
items: WorkerAutoImportBatchItem[],
requestId: string,
concurrency: number,
sessionGapThreshold?: number,
signal?: AbortSignal
): Promise<AutoImportBatchItemResult[]> {
return sharedAutoImportBatch(
items.map((item) => ({
id: item.id,
filePath: item.filePath,
options: {
explicitSessionId: item.explicitSessionId,
formatOptions: item.formatOptions,
},
})),
{
listSessionIds: () =>
fs.existsSync(getDbDir())
? fs
.readdirSync(getDbDir())
.filter((name) => name.endsWith('.db'))
.map((name) => name.slice(0, -3))
: [],
openReadonly: (sessionId) => new BetterSqliteAdapter(openRawDatabase(getDbPath(sessionId), { readonly: true })),
sessionExists: (sessionId) => fs.existsSync(getDbPath(sessionId)),
createSession: (sourcePath, sourceFormatOptions, sessionId) =>
streamImport(sourcePath, requestId, sourceFormatOptions, sessionId, sessionGapThreshold),
appendSession: (sessionId, sourcePath, sourceFormatOptions) =>
incrementalImport(sessionId, sourcePath, requestId, {
formatId: typeof sourceFormatOptions?.formatId === 'string' ? sourceFormatOptions.formatId : undefined,
chatIndex: typeof sourceFormatOptions?.chatIndex === 'number' ? sourceFormatOptions.chatIndex : undefined,
}),
},
{
concurrency,
signal,
onItemStart: (_item, index) => {
sendProgress(requestId, {
stage: 'detecting',
percentage: 0,
message: '',
bytesRead: 0,
totalBytes: 0,
messagesProcessed: 0,
batchIndex: index,
batchEvent: 'start',
} as Parameters<typeof sendProgress>[1])
},
onItemProgress: (_item, index, progress) => {
sendProgress(requestId, {
...progress,
batchIndex: index,
batchEvent: 'progress',
} as Parameters<typeof sendProgress>[1])
},
onItemComplete: (_item, index, result) => {
sendProgress(requestId, {
stage: 'done',
percentage: 100,
message: '',
bytesRead: 0,
totalBytes: 0,
messagesProcessed: 0,
batchIndex: index,
batchEvent: 'complete',
batchResult: result,
} as Parameters<typeof sendProgress>[1])
},
}
)
}
/**
* Dry-run analysis: parse without writing to DB.
*/
+96 -11
View File
@@ -18,9 +18,15 @@ import {
IMPORT_IN_PROGRESS_ERROR_KEY,
ImportInProgressError,
raiseChatDbCompatibilityGate,
resolveDefaultBatchConcurrency,
withDataDirImportLock,
} from '@openchatlab/node-runtime'
import type { PushImportAnalysisOutcome, PushImportOutcome, PushImportPayload } from '@openchatlab/node-runtime'
import type {
AutoImportBatchItemResult,
PushImportAnalysisOutcome,
PushImportOutcome,
PushImportPayload,
} from '@openchatlab/node-runtime'
import { getWorkerRequestTimeoutMs, isRestartableReadOnlyRequestType } from './workerTimeoutPolicy'
interface WorkerRequestOptions {
@@ -37,7 +43,7 @@ const pendingRequests = new Map<
{
resolve: (value: any) => void
reject: (error: Error) => void
timeout: ReturnType<typeof setTimeout>
timeout?: ReturnType<typeof setTimeout>
restartOnTimeout: boolean
onProgress?: (progress: ParseProgress) => void // 进度回调
}
@@ -55,7 +61,7 @@ function hasProgressRequest(): boolean {
function rejectAllPending(error: Error): void {
for (const [id, pending] of pendingRequests.entries()) {
clearTimeout(pending.timeout)
if (pending.timeout) clearTimeout(pending.timeout)
pending.reject(error)
pendingRequests.delete(id)
}
@@ -163,7 +169,7 @@ export function initWorker(): void {
// 处理完成或错误消息
pendingRequests.delete(id)
clearTimeout(pending.timeout)
if (pending.timeout) clearTimeout(pending.timeout)
if (success) {
pending.resolve(result)
@@ -243,7 +249,7 @@ function sendToWorkerWithProgress<T>(
type: string,
payload: any,
onProgress?: (progress: ParseProgress) => void,
timeoutMs: number = 600000 // 默认 10 分钟超时
timeoutMs: number | null = 600000 // 默认 10 分钟超时;null 表示由调用方持有直至 Worker 完成
): Promise<T> {
return new Promise((resolve, reject) => {
if (!worker) {
@@ -258,12 +264,15 @@ function sendToWorkerWithProgress<T>(
const requestWorker = worker!
const id = `req_${++requestIdCounter}`
const timeout = setTimeout(() => {
if (pendingRequests.has(id)) {
pendingRequests.delete(id)
reject(new Error(`Worker request timeout: ${type}`))
}
}, timeoutMs)
const timeout =
timeoutMs === null
? undefined
: setTimeout(() => {
if (pendingRequests.has(id)) {
pendingRequests.delete(id)
reject(new Error(`Worker request timeout: ${type}`))
}
}, timeoutMs)
pendingRequests.set(id, { resolve, reject, timeout, restartOnTimeout: false, onProgress })
@@ -619,6 +628,82 @@ export async function autoImport(
}
}
export interface DesktopAutoImportBatchItem {
id: string
filePath: string
formatOptions?: Record<string, unknown>
explicitSessionId?: string
}
export interface DesktopAutoImportBatchProgress extends ParseProgress {
batchIndex: number
batchEvent: 'start' | 'progress' | 'complete'
batchResult?: AutoImportBatchItemResult
}
export async function autoImportBatch(
batchId: string,
items: DesktopAutoImportBatchItem[],
onProgress?: (progress: DesktopAutoImportBatchProgress) => void,
options?: { concurrency?: number; sessionGapThreshold?: number }
): Promise<AutoImportBatchItemResult[]> {
try {
return await withDataDirImportLock(getPathProvider().getUserDataDir(), async () => {
assertDataDirCompatibleNow()
const results = await sendToWorkerWithProgress<AutoImportBatchItemResult[]>(
'autoImportBatch',
{
batchId,
items,
concurrency: options?.concurrency ?? resolveDefaultBatchConcurrency(items.length),
sessionGapThreshold: options?.sessionGapThreshold,
},
onProgress as (progress: ParseProgress) => void,
null
)
if (!results.some((result) => result.status === 'success')) return results
try {
raiseChatDbCompatibilityGate(getPathProvider(), {
version: getDesktopAppVersion(app.getVersion()),
kind: 'desktop',
})
} catch (error) {
const message = error instanceof Error ? error.message : String(error)
for (const result of results) {
if (result.status === 'success' && result.result.importMode === 'created' && result.result.sessionId) {
deleteImportedSession(result.result.sessionId)
}
}
return results.map((result) =>
result.status === 'success'
? {
id: result.id,
status: 'failed' as const,
error: message,
result: { ...result.result, success: false, error: message },
}
: result
)
}
return results
})
} catch (error) {
if (error instanceof ImportInProgressError) {
return items.map((item) => ({
id: item.id,
status: 'failed',
error: IMPORT_IN_PROGRESS_ERROR_KEY,
}))
}
throw error
}
}
export async function cancelAutoImportBatch(batchId: string): Promise<void> {
await sendToWorker('cancelAutoImportBatch', { batchId })
}
/**
* 在 worker 中执行共享 Push Import 写库语义;主进程持有全局导入锁,
* 并在 worker 成功后完成数据目录兼容门禁更新。
+14
View File
@@ -59,6 +59,20 @@ export const chatApi = {
importWithOptions: (filePath: string, formatOptions: Record<string, unknown>): Promise<ChatImportResult> =>
ipcRenderer.invoke('chat:importWithOptions', filePath, formatOptions),
importBatch: (
batchId: string,
items: Array<{ id: string; filePath: string }>,
options?: { concurrency?: number; sessionGapThreshold?: number }
) => ipcRenderer.invoke('chat:importBatch', batchId, items, options),
cancelImportBatch: (batchId: string) => ipcRenderer.invoke('chat:cancelImportBatch', batchId),
onImportBatchProgress: (callback: (progress: Record<string, unknown>) => void) => {
const handler = (_event: Electron.IpcRendererEvent, progress: Record<string, unknown>) => callback(progress)
ipcRenderer.on('chat:importBatchProgress', handler)
return () => ipcRenderer.removeListener('chat:importBatchProgress', handler)
},
scanMultiChatFile: (
filePath: string
): Promise<{
+24
View File
@@ -88,6 +88,30 @@ interface ChatApi {
import: (filePath: string) => Promise<ChatImportResult>
importDirectory: (dirPath: string, options?: Record<string, unknown>) => Promise<ChatImportResult>
importWithOptions: (filePath: string, formatOptions: Record<string, unknown>) => Promise<ChatImportResult>
importBatch: (
batchId: string,
items: Array<{ id: string; filePath: string }>,
options?: { concurrency?: number; sessionGapThreshold?: number }
) => Promise<
Array<{
id: string
status: 'success' | 'failed' | 'cancelled'
result?: ChatImportResult
error?: string
}>
>
cancelImportBatch: (batchId: string) => Promise<{ success: boolean }>
onImportBatchProgress: (
callback: (progress: {
batchId: string
batchIndex: number
batchEvent: 'start' | 'progress' | 'complete'
stage: string
percentage: number
message?: string
batchResult?: unknown
}) => void
) => () => void
scanMultiChatFile: (filePath: string) => Promise<{
success: boolean
chats: Array<{ index: number; name: string; type: string; id: number; messageCount: number }>