🏄‍♂️ refactor: Optimize Reasoning UI & Token Streaming (#5546)

* ✨ feat: Implement Show Thinking feature; refactor: testing thinking render optimizations

* ✨ feat: Refactor Thinking component styles and enhance Markdown rendering

* chore: add back removed code, revert type changes

* chore: Add back resetCounter effect to Markdown component for improved code block indexing

* chore: bump @librechat/agents and google langchain packages

* WIP: reasoning type updates

* WIP: first pass, reasoning content blocks

* chore: revert code

* chore: bump @librechat/agents

* refactor: optimize reasoning tag handling

* style: ul indent padding

* feat: add Reasoning component to handle reasoning display

* feat: first pass, content reasoning part styling

* refactor: add content placeholder for endpoints using new stream handler

* refactor: only cache messages when requesting stream audio

* fix: circular dep.

* fix: add default param

* refactor: tts, only request after message stream, fix chrome autoplay

* style: update label for submitting state and add localization for 'Thinking...'

* fix: improve global audio pause logic and reset active run ID

* fix: handle artifact edge cases

* fix: remove unnecessary console log from artifact update test

* feat: add support for continued message handling with new streaming method

---------

Co-authored-by: Marco Beretta <81851188+berry-13@users.noreply.github.com>

api/server/services/Artifacts/update.js

@@ -57,14 +57,42 @@ const findAllArtifacts = (message) => {
 
 const replaceArtifactContent = (originalText, artifact, original, updated) => {
   const artifactContent = artifact.text.substring(artifact.start, artifact.end);
-  const relativeIndex = artifactContent.indexOf(original);
+
+  // Find boundaries between ARTIFACT_START and ARTIFACT_END
+  const contentStart = artifactContent.indexOf('\n', artifactContent.indexOf(ARTIFACT_START)) + 1;
+  const contentEnd = artifactContent.lastIndexOf(ARTIFACT_END);
+
+  if (contentStart === -1 || contentEnd === -1) {
+    return null;
+  }
+
+  // Check if there are code blocks
+  const codeBlockStart = artifactContent.indexOf('```\n', contentStart);
+  const codeBlockEnd = artifactContent.lastIndexOf('\n```', contentEnd);
+
+  // Determine where to look for the original content
+  let searchStart, searchEnd;
+  if (codeBlockStart !== -1 && codeBlockEnd !== -1) {
+    // If code blocks exist, search between them
+    searchStart = codeBlockStart + 4; // after ```\n
+    searchEnd = codeBlockEnd;
+  } else {
+    // Otherwise search in the whole artifact content
+    searchStart = contentStart;
+    searchEnd = contentEnd;
+  }
+
+  const innerContent = artifactContent.substring(searchStart, searchEnd);
+  // Remove trailing newline from original for comparison
+  const originalTrimmed = original.replace(/\n$/, '');
+  const relativeIndex = innerContent.indexOf(originalTrimmed);
 
   if (relativeIndex === -1) {
     return null;
   }
 
-  const absoluteIndex = artifact.start + relativeIndex;
-  const endText = originalText.substring(absoluteIndex + original.length);
+  const absoluteIndex = artifact.start + searchStart + relativeIndex;
+  const endText = originalText.substring(absoluteIndex + originalTrimmed.length);
   const hasTrailingNewline = endText.startsWith('\n');
 
   const updatedText =

api/server/services/Artifacts/update.spec.js

@@ -260,8 +260,61 @@ console.log(greeting);`;
       codeExample,
       'updated content',
     );
-    console.log(result);
     expect(result).toMatch(/id="2".*updated content/s);
     expect(result).toMatch(new RegExp(`${ARTIFACT_START}.*updated content.*${ARTIFACT_END}`, 's'));
   });
+
+  test('should handle empty content in artifact without code blocks', () => {
+    const artifactText = `${ARTIFACT_START}\n\n${ARTIFACT_END}`;
+    const artifact = {
+      start: 0,
+      end: artifactText.length,
+      text: artifactText,
+      source: 'text',
+    };
+
+    const result = replaceArtifactContent(artifactText, artifact, '', 'new content');
+    expect(result).toBe(`${ARTIFACT_START}\nnew content\n${ARTIFACT_END}`);
+  });
+
+  test('should handle empty content in artifact with code blocks', () => {
+    const artifactText = createArtifactText({ content: '' });
+    const artifact = {
+      start: 0,
+      end: artifactText.length,
+      text: artifactText,
+      source: 'text',
+    };
+
+    const result = replaceArtifactContent(artifactText, artifact, '', 'new content');
+    expect(result).toMatch(/```\nnew content\n```/);
+  });
+
+  test('should handle content with trailing newline in code blocks', () => {
+    const contentWithNewline = 'console.log("test")\n';
+    const message = {
+      text: `Some prefix text\n${createArtifactText({
+        content: contentWithNewline,
+      })}\nSome suffix text`,
+    };
+
+    const artifacts = findAllArtifacts(message);
+    expect(artifacts).toHaveLength(1);
+
+    const result = replaceArtifactContent(
+      message.text,
+      artifacts[0],
+      contentWithNewline,
+      'updated content',
+    );
+
+    // Should update the content and preserve artifact structure
+    expect(result).toContain('```\nupdated content\n```');
+    // Should preserve surrounding text
+    expect(result).toMatch(/^Some prefix text\n/);
+    expect(result).toMatch(/\nSome suffix text$/);
+    // Should not have extra newlines
+    expect(result).not.toContain('\n\n```');
+    expect(result).not.toContain('```\n\n');
+  });
 });

api/server/services/Files/Audio/TTSService.js

@@ -364,7 +364,7 @@ class TTSService {
       shouldContinue = false;
     });
 
-    const processChunks = createChunkProcessor(req.body.messageId);
+    const processChunks = createChunkProcessor(req.user.id, req.body.messageId);
 
     try {
       while (shouldContinue) {

api/server/services/Files/Audio/streamAudio.js

@@ -1,4 +1,5 @@
-const { CacheKeys, findLastSeparatorIndex, SEPARATORS } = require('librechat-data-provider');
+const { CacheKeys, findLastSeparatorIndex, SEPARATORS, Time } = require('librechat-data-provider');
+const { getMessage } = require('~/models/Message');
 const { getLogStores } = require('~/cache');
 
 /**
@@ -47,10 +48,11 @@ const MAX_NOT_FOUND_COUNT = 6;
 const MAX_NO_CHANGE_COUNT = 10;
 
 /**
+ * @param {string} user
  * @param {string} messageId
  * @returns {() => Promise<{ text: string, isFinished: boolean }[]>}
  */
-function createChunkProcessor(messageId) {
+function createChunkProcessor(user, messageId) {
   let notFoundCount = 0;
   let noChangeCount = 0;
   let processedText = '';
@@ -73,15 +75,27 @@ function createChunkProcessor(messageId) {
     }
 
     /** @type { string | { text: string; complete: boolean } } */
-    const message = await messageCache.get(messageId);
+    let message = await messageCache.get(messageId);
+    if (!message) {
+      message = await getMessage({ user, messageId });
+    }
 
     if (!message) {
       notFoundCount++;
       return [];
+    } else {
+      messageCache.set(
+        messageId,
+        {
+          text: message.text,
+          complete: true,
+        },
+        Time.FIVE_MINUTES,
+      );
     }
 
     const text = typeof message === 'string' ? message : message.text;
-    const complete = typeof message === 'string' ? false : message.complete;
+    const complete = typeof message === 'string' ? false : message.complete ?? true;
 
     if (text === processedText) {
       noChangeCount++;

api/server/services/Files/Audio/streamAudio.spec.js

@@ -3,6 +3,13 @@ const { createChunkProcessor, splitTextIntoChunks } = require('./streamAudio');
 jest.mock('keyv');
 
 const globalCache = {};
+jest.mock('~/models/Message', () => {
+  return {
+    getMessage: jest.fn().mockImplementation((messageId) => {
+      return globalCache[messageId] || null;
+    }),
+  };
+});
 jest.mock('~/cache/getLogStores', () => {
   return jest.fn().mockImplementation(() => {
     const EventEmitter = require('events');
@@ -56,9 +63,10 @@ describe('processChunks', () => {
     jest.resetAllMocks();
     mockMessageCache = {
       get: jest.fn(),
+      set: jest.fn(),
     };
     require('~/cache/getLogStores').mockReturnValue(mockMessageCache);
-    processChunks = createChunkProcessor('message-id');
+    processChunks = createChunkProcessor('userId', 'message-id');
   });
 
   it('should return an empty array when the message is not found', async () => {

api/server/services/Runs/StreamRunManager.js

@@ -1,19 +1,15 @@
-const throttle = require('lodash/throttle');
 const {
-  Time,
-  CacheKeys,
+  Constants,
   StepTypes,
   ContentTypes,
   ToolCallTypes,
   MessageContentTypes,
   AssistantStreamEvents,
-  Constants,
 } = require('librechat-data-provider');
 const { retrieveAndProcessFile } = require('~/server/services/Files/process');
 const { processRequiredActions } = require('~/server/services/ToolService');
 const { createOnProgress, sendMessage, sleep } = require('~/server/utils');
 const { processMessages } = require('~/server/services/Threads');
-const { getLogStores } = require('~/cache');
 const { logger } = require('~/config');
 
 /**
@@ -611,20 +607,8 @@ class StreamRunManager {
       const index = this.getStepIndex(stepKey);
       this.orderedRunSteps.set(index, message_creation);
 
-      const messageCache = getLogStores(CacheKeys.MESSAGES);
-      // Create the Factory Function to stream the message
-      const { onProgress: progressCallback } = createOnProgress({
-        onProgress: throttle(
-          () => {
-            messageCache.set(this.finalMessage.messageId, this.getText(), Time.FIVE_MINUTES);
-          },
-          3000,
-          { trailing: false },
-        ),
-      });
+      const { onProgress: progressCallback } = createOnProgress();
 
-      // This creates a function that attaches all of the parameters
-      // specified here to each SSE message generated by the TextStream
       const onProgress = progressCallback({
         index,
         res: this.res,