🧠 refactor: Memory Timeout after Completion and Guarantee Stream Final Event (#8955)

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<string | undefined>}
    */
@@ -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',

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) {