diff --git a/packages/backend/server/src/plugins/copilot/context/job.ts b/packages/backend/server/src/plugins/copilot/context/job.ts index a44b7cc69f..1b71eaa8bf 100644 --- a/packages/backend/server/src/plugins/copilot/context/job.ts +++ b/packages/backend/server/src/plugins/copilot/context/job.ts @@ -73,11 +73,18 @@ export class CopilotContextDocJob { } @OnEvent('workspace.doc.embedding') - async addDocEmbeddingQueue(docs: Events['workspace.doc.embedding']) { + async addDocEmbeddingQueue( + docs: Events['workspace.doc.embedding'], + contextId?: string + ) { if (!this.supportEmbedding) return; for (const { workspaceId, docId } of docs) { - await this.queue.add('doc.embedPendingDocs', { workspaceId, docId }); + await this.queue.add('doc.embedPendingDocs', { + contextId, + workspaceId, + docId, + }); } } @@ -130,23 +137,24 @@ export class CopilotContextDocJob { fileId, chunkSize: total, }); - } catch (e: any) { - const error = mapAnyError(e); - error.log('CopilotJob', { - workspaceId, - fileId, - }); - + } catch (error: any) { this.event.emit('workspace.file.embed.failed', { contextId, fileId, - error: e.toString(), + error: mapAnyError(error).message, }); + + // passthrough error to job queue + throw error; } } @OnJob('doc.embedPendingDocs') - async embedPendingDocs({ workspaceId, docId }: Jobs['doc.embedPendingDocs']) { + async embedPendingDocs({ + contextId, + workspaceId, + docId, + }: Jobs['doc.embedPendingDocs']) { if (!this.supportEmbedding) return; try { @@ -165,11 +173,16 @@ export class CopilotContextDocJob { ); } } - } catch (e: any) { - this.logger.error( - `Failed to embed pending doc: ${workspaceId}::${docId}`, - e - ); + } catch (error: any) { + if (contextId) { + this.event.emit('workspace.doc.embed.failed', { + contextId, + docId, + }); + } + + // passthrough error to job queue + throw error; } } } diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts index e12da3a5ea..7117cb483e 100644 --- a/packages/backend/server/src/plugins/copilot/context/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts @@ -464,7 +464,8 @@ export class CopilotContextResolver { options.docs.map(docId => ({ workspaceId: session.workspaceId, docId, - })) + })), + session.id ); } @@ -523,12 +524,10 @@ export class CopilotContextResolver { try { const record = await session.addDocRecord(options.docId); - await this.jobs.addDocEmbeddingQueue([ - { - workspaceId: session.workspaceId, - docId: options.docId, - }, - ]); + await this.jobs.addDocEmbeddingQueue( + [{ workspaceId: session.workspaceId, docId: options.docId }], + session.id + ); return { ...record, status: record.status || null }; } catch (e: any) { diff --git a/packages/backend/server/src/plugins/copilot/context/service.ts b/packages/backend/server/src/plugins/copilot/context/service.ts index b9b5bf2307..f666733163 100644 --- a/packages/backend/server/src/plugins/copilot/context/service.ts +++ b/packages/backend/server/src/plugins/copilot/context/service.ts @@ -11,6 +11,7 @@ import { import { ContextConfig, ContextConfigSchema, + ContextDoc, ContextEmbedStatus, ContextFile, Models, @@ -148,6 +149,18 @@ export class CopilotContextService implements OnApplicationBootstrap { return null; } + @OnEvent('workspace.doc.embed.failed') + async onDocEmbedFailed({ + contextId, + docId, + }: Events['workspace.doc.embed.failed']) { + const context = await this.get(contextId); + await context.saveDocRecord(docId, doc => ({ + ...(doc as ContextDoc), + status: ContextEmbedStatus.failed, + })); + } + @OnEvent('workspace.file.embed.finished') async onFileEmbedFinish({ contextId, diff --git a/packages/backend/server/src/plugins/copilot/context/session.ts b/packages/backend/server/src/plugins/copilot/context/session.ts index 39de23cef4..34047df568 100644 --- a/packages/backend/server/src/plugins/copilot/context/session.ts +++ b/packages/backend/server/src/plugins/copilot/context/session.ts @@ -217,6 +217,23 @@ export class ContextSession implements AsyncDisposable { ); } + async saveDocRecord( + docId: string, + cb: ( + record: Pick & + Partial> + ) => ContextDoc + ) { + const docs = [this.config.docs, ...this.config.categories.map(c => c.docs)] + .flat() + .filter(d => d.id === docId); + for (const doc of docs) { + Object.assign(doc, cb({ ...doc })); + } + + await this.save(); + } + async saveFileRecord( fileId: string, cb: ( diff --git a/packages/backend/server/src/plugins/copilot/context/types.ts b/packages/backend/server/src/plugins/copilot/context/types.ts index fbd6a6a50e..ac089b9959 100644 --- a/packages/backend/server/src/plugins/copilot/context/types.ts +++ b/packages/backend/server/src/plugins/copilot/context/types.ts @@ -10,11 +10,18 @@ declare global { workspaceId: string; docId: string; }>; + + 'workspace.doc.embed.failed': { + contextId: string; + docId: string; + }; + 'workspace.file.embed.finished': { contextId: string; fileId: string; chunkSize: number; }; + 'workspace.file.embed.failed': { contextId: string; fileId: string; @@ -23,6 +30,7 @@ declare global { } interface Jobs { 'doc.embedPendingDocs': { + contextId?: string; workspaceId: string; docId: string; };