diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 4f4b84d534..25d98aa3fc 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -402,6 +402,34 @@ class AgentClient extends BaseClient { return result; } + /** + * Creates a promise that resolves with the memory promise result or undefined after a timeout + * @param {Promise<(TAttachment | null)[] | undefined>} memoryPromise - The memory promise to await + * @param {number} timeoutMs - Timeout in milliseconds (default: 3000) + * @returns {Promise<(TAttachment | null)[] | undefined>} + */ + async awaitMemoryWithTimeout(memoryPromise, timeoutMs = 3000) { + if (!memoryPromise) { + return; + } + + try { + const timeoutPromise = new Promise((_, reject) => + setTimeout(() => reject(new Error('Memory processing timeout')), timeoutMs), + ); + + const attachments = await Promise.race([memoryPromise, timeoutPromise]); + return attachments; + } catch (error) { + if (error.message === 'Memory processing timeout') { + logger.warn('[AgentClient] Memory processing timed out after 3 seconds'); + } else { + logger.error('[AgentClient] Error processing memory:', error); + } + return; + } + } + /** * @returns {Promise} */ @@ -1002,11 +1030,9 @@ class AgentClient extends BaseClient { }); try { - if (memoryPromise) { - const attachments = await memoryPromise; - if (attachments && attachments.length > 0) { - this.artifactPromises.push(...attachments); - } + const attachments = await this.awaitMemoryWithTimeout(memoryPromise); + if (attachments && attachments.length > 0) { + this.artifactPromises.push(...attachments); } await this.recordCollectedUsage({ context: 'message' }); } catch (err) { @@ -1016,11 +1042,9 @@ class AgentClient extends BaseClient { ); } } catch (err) { - if (memoryPromise) { - const attachments = await memoryPromise; - if (attachments && attachments.length > 0) { - this.artifactPromises.push(...attachments); - } + const attachments = await this.awaitMemoryWithTimeout(memoryPromise); + if (attachments && attachments.length > 0) { + this.artifactPromises.push(...attachments); } logger.error( '[api/server/controllers/agents/client.js #sendCompletion] Operation aborted', diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index b4a0e40b24..8054a6a680 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -233,6 +233,26 @@ const AgentController = async (req, res, next, initializeClient, addTitle) => { ); } } + // Edge case: sendMessage completed but abort happened during sendCompletion + // We need to ensure a final event is sent + else if (!res.headersSent && !res.finished) { + logger.debug( + '[AgentController] Handling edge case: `sendMessage` completed but aborted during `sendCompletion`', + ); + + const finalResponse = { ...response }; + finalResponse.error = true; + + sendEvent(res, { + final: true, + conversation, + title: conversation.title, + requestMessage: userMessage, + responseMessage: finalResponse, + error: { message: 'Request was aborted during completion' }, + }); + res.end(); + } // Save user message if needed if (!client.skipSaveUserMessage) {