Fix thought and final response in the same content, thought won't render

src/app/components/chat/chat.component.spec.ts

@@ -1401,17 +1401,20 @@ describe('ChatComponent', () => {
       expect(uiEvent.text).toBe('hello world');
     });
 
-    it('should extract and parse inline <a2ui-json> block, and strip it from text', () => {
+    it('should extract and parse inline <a2ui-json> block, and strip it from text and textParts', () => {
       const payload = [{beginRendering: {surfaceId: 'cloud_dash'}}];
+      const text = `Here is the UI:\n<a2ui-json>\n${JSON.stringify(payload)}\n</a2ui-json>\nEnjoy!`;
       const uiEvent = new UiEvent({
         role: 'bot',
-        text: `Here is the UI:\n<a2ui-json>\n${JSON.stringify(payload)}\n</a2ui-json>\nEnjoy!`,
+        text,
+        textParts: [{ text, thought: false }],
         event: {} as any
       });
 
       (component as any).extractA2uiJsonFromText(uiEvent);
       expect(uiEvent.a2uiData).toEqual({beginRendering: {beginRendering: {surfaceId: 'cloud_dash'}}});
       expect(uiEvent.text).toBe('Here is the UI:\n\nEnjoy!');
+      expect(uiEvent.textParts).toEqual([{ text: 'Here is the UI:\n\nEnjoy!', thought: false }]);
     });
 
     it('should keep tags and log warning if JSON parsing fails', () => {
@@ -1427,6 +1430,75 @@ describe('ChatComponent', () => {
     });
   });
 
+  describe('thought and text parts processing', () => {
+    it('should build UiEvent with textParts separating thoughts and final response text', () => {
+      const event = {
+        id: 'event-thought-test',
+        author: 'bot',
+        content: {
+          parts: [
+            {
+              text: '/*PLANNING*/Thinking about the question...',
+              thought: true
+            },
+            {
+              text: 'The actual response text.',
+            }
+          ]
+        }
+      };
+
+      const uiEvent = (component as any).buildUiEventFromEvent(event);
+
+      expect(uiEvent.text).toBe('Thinking about the question...The actual response text.');
+      expect(uiEvent.textParts).toEqual([
+        { text: 'Thinking about the question...', thought: true },
+        { text: 'The actual response text.', thought: false }
+      ]);
+    });
+
+    it('should correctly merge streaming partial updates with thought and non-thought parts', () => {
+      const initialEvent = new UiEvent({
+        role: 'bot',
+        text: 'Thinking',
+        thought: true,
+        textParts: [{ text: 'Thinking', thought: true }],
+        event: { id: 'stream-event', partial: true } as any
+      });
+
+      // Stream update 1: more thought text
+      const update1 = {
+        id: 'stream-event',
+        partial: true,
+        content: {
+          parts: [{ text: ' further...', thought: true }]
+        }
+      };
+
+      let merged = (component as any).mergePartialEvent(initialEvent, update1);
+      expect(merged.text).toBe('Thinking further...');
+      expect(merged.textParts).toEqual([
+        { text: 'Thinking further...', thought: true }
+      ]);
+
+      // Stream update 2: final answer text
+      const update2 = {
+        id: 'stream-event',
+        partial: true,
+        content: {
+          parts: [{ text: 'Here is the answer!' }]
+        }
+      };
+
+      merged = (component as any).mergePartialEvent(merged, update2);
+      expect(merged.text).toBe('Thinking further...Here is the answer!');
+      expect(merged.textParts).toEqual([
+        { text: 'Thinking further...', thought: true },
+        { text: 'Here is the answer!', thought: false }
+      ]);
+    });
+  });
+
   describe('refreshLatestSession', () => {
     beforeEach(() => {
       mockAgentService.listAppsResponse.next([TEST_APP_1_NAME]);

src/app/components/chat/chat.component.ts

@@ -1156,6 +1156,7 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
 
   submitAgentRunRequest(req: AgentRunRequest) {
     this.autoSelectLatestEvent = true;
+
     this.activeSseSubscription = this.agentService.runSse(req).subscribe({
       next: async (chunkJson: any) => {
         if (chunkJson.error) {
@@ -1362,7 +1363,11 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
   }
 
   private mergePartialEvent(lastEvent: UiEvent, apiEvent: any): UiEvent {
-    const updatedEvent = new UiEvent({ ...lastEvent, event: apiEvent as any });
+    const updatedEvent = new UiEvent({
+      ...lastEvent,
+      event: apiEvent as any,
+      textParts: lastEvent.textParts ? lastEvent.textParts.map(p => ({ ...p })) : undefined
+    });
 
     let parts = apiEvent.content?.parts || [];
     if (this.isEventA2aResponse(apiEvent)) {
@@ -1372,15 +1377,15 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
 
     parts.forEach((part: any) => {
       if (part.text !== undefined && part.text !== null) {
-        updatedEvent.text = (updatedEvent.text || '') + part.text;
-        if (part.thought) {
-          updatedEvent.thought = true;
-          updatedEvent.text = this.processThoughtText(updatedEvent.text || '');
-        }
+        const processedText = part.thought ? this.processThoughtText(part.text) : part.text;
+        updatedEvent.text = (updatedEvent.text || '') + processedText;
+        const isThought = !!part.thought;
+        this.addTextToParts(updatedEvent, processedText, isThought);
       } else {
         this.processPartIntoMessage(part, apiEvent, updatedEvent);
       }
     });
+    updatedEvent.thought = updatedEvent.textParts?.every(p => p.thought) ?? false;
 
     if (apiEvent.inputTranscription) {
       const previousText = (lastEvent.event as any)?.inputTranscription?.text || '';
@@ -1431,12 +1436,13 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     let combinedTextPart: Part | undefined;
 
     for (const part of parts) {
-      if (part.text && !part.thought) {
-        if (!combinedTextPart) {
-          combinedTextPart = { text: part.text };
-          result.push(combinedTextPart);
-        } else {
+      if (part.text) {
+        const isThought = !!part.thought;
+        if (combinedTextPart && combinedTextPart.text && !!combinedTextPart.thought === isThought) {
           combinedTextPart.text += part.text;
+        } else {
+          combinedTextPart = { text: part.text, thought: isThought };
+          result.push(combinedTextPart);
         }
       } else {
         combinedTextPart = undefined;
@@ -1578,6 +1584,23 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
       const afterText = uiEvent.text.substring(endIndex + endTag.length);
       uiEvent.text = (beforeText + afterText).trim();
 
+      // Strip from uiEvent.textParts as well
+      if (uiEvent.textParts) {
+        for (const part of uiEvent.textParts) {
+          const partStartIndex = part.text.indexOf(startTag);
+          if (partStartIndex !== -1) {
+            const partEndIndex = part.text.indexOf(endTag, partStartIndex + startTag.length);
+            if (partEndIndex !== -1) {
+              const partBefore = part.text.substring(0, partStartIndex);
+              const partAfter = part.text.substring(partEndIndex + endTag.length);
+              part.text = (partBefore + partAfter).trim();
+            }
+          }
+        }
+        // Remove any parts that became empty after stripping
+        uiEvent.textParts = uiEvent.textParts.filter(part => part.text.trim().length > 0);
+      }
+
     } catch (e) {
       console.warn('Failed to parse inline <a2ui-json> block from text:', e);
     }
@@ -1600,6 +1623,19 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     return `data:${mimeType};base64,${fixedBase64Data}`;
   }
 
+  private addTextToParts(uiEvent: UiEvent, text: string, thought: boolean) {
+    if (!text) return;
+    if (!uiEvent.textParts) {
+      uiEvent.textParts = [];
+    }
+    const lastPart = uiEvent.textParts[uiEvent.textParts.length - 1];
+    if (lastPart && !!lastPart.thought === thought) {
+      lastPart.text += text;
+    } else {
+      uiEvent.textParts.push({ text, thought });
+    }
+  }
+
   private processPartIntoMessage(part: any, event: any, uiEvent: UiEvent) {
     if (!part) return;
 
@@ -1617,8 +1653,10 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     }
 
     if (part.text) {
-      uiEvent.text = (uiEvent.text || '') + part.text;
-      uiEvent.thought = part.thought ? true : false;
+      const processedText = part.thought ? this.processThoughtText(part.text) : part.text;
+      uiEvent.text = (uiEvent.text || '') + processedText;
+      this.addTextToParts(uiEvent, processedText, !!part.thought);
+      uiEvent.thought = uiEvent.textParts?.every(p => p.thought) ?? false;
       if (event?.groundingMetadata && event.groundingMetadata.searchEntryPoint &&
         event.groundingMetadata.searchEntryPoint.renderedContent) {
         uiEvent.renderedContent =
@@ -2213,9 +2251,10 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
 
   private buildUiEventFromEvent(event: any, reverseOrder: boolean = false): UiEvent {
     const isA2aResponse = this.isEventA2aResponse(event);
-    const parts = isA2aResponse ?
+    const rawParts = isA2aResponse ?
       this.combineA2uiDataParts(event.content?.parts) :
       event.content?.parts || [];
+    const parts = this.combineTextParts(rawParts);
     const partsToProcess = reverseOrder ? [...parts].reverse() : parts;
 
     const role = event.author === 'user' ? 'user' : 'bot';

src/app/components/content-bubble/content-bubble.component.html

@@ -27,7 +27,18 @@
     }
     @if (uiEvent.event.nodeInfo?.['messageAsOutput']) {
       <div [appJsonTooltip]="jsonOutputData" appJsonTooltipTitle="Node Output">
-        <ng-container [ngComponentOutlet]="markdownComponent" [ngComponentOutletInputs]="{text: uiEvent.text || rawMessageText, thought: uiEvent.thought}"></ng-container>
+        @if (uiEvent.textParts && uiEvent.textParts.length > 0) {
+          @for (part of uiEvent.textParts; track part; let first = $first) {
+            <div [class.thought-container]="part.thought && type !== 'thought'" [class.not-first-part]="!first">
+              @if (part.thought && type !== 'thought') {
+                <div class="output-chip-header">thought</div>
+              }
+              <ng-container [ngComponentOutlet]="markdownComponent" [ngComponentOutletInputs]="{text: part.text, thought: part.thought}"></ng-container>
+            </div>
+          }
+        } @else {
+          <ng-container [ngComponentOutlet]="markdownComponent" [ngComponentOutletInputs]="{text: uiEvent.text || rawMessageText, thought: uiEvent.thought}"></ng-container>
+        }
       </div>
     } @else if (uiEvent.thought || uiEvent.text || uiEvent.renderedContent || uiEvent.a2uiData || uiEvent.event.inputTranscription || uiEvent.event.outputTranscription) {
       <div>
@@ -62,7 +73,18 @@
             </div>
           </div>
           } @else {
-            <ng-container [ngComponentOutlet]="markdownComponent" [ngComponentOutletInputs]="{text: uiEvent.text!, thought: uiEvent.thought}"></ng-container>
+            @if (uiEvent.textParts && uiEvent.textParts.length > 0) {
+              @for (part of uiEvent.textParts; track part; let first = $first) {
+                <div [class.thought-container]="part.thought && type !== 'thought'" [class.not-first-part]="!first">
+                  @if (part.thought && type !== 'thought') {
+                    <div class="output-chip-header">thought</div>
+                  }
+                  <ng-container [ngComponentOutlet]="markdownComponent" [ngComponentOutletInputs]="{text: part.text, thought: part.thought}"></ng-container>
+                </div>
+              }
+            } @else {
+              <ng-container [ngComponentOutlet]="markdownComponent" [ngComponentOutletInputs]="{text: uiEvent.text!, thought: uiEvent.thought}"></ng-container>
+            }
           }
         }
         </div>

src/app/components/content-bubble/content-bubble.component.scss

@@ -120,6 +120,22 @@
   }
 }
 
+.thought-container {
+  margin-bottom: 12px;
+  padding-bottom: 8px;
+  border-bottom: 1px dashed var(--mat-sys-outline-variant, rgba(0, 0, 0, 0.1));
+
+  &:last-child {
+    margin-bottom: 0;
+    padding-bottom: 0;
+    border-bottom: none;
+  }
+}
+
+.not-first-part {
+  margin-top: 12px;
+}
+
 @media (max-width: 768px) {
   .content-bubble {
     padding: 5px 12px !important;

src/app/core/models/UiEvent.ts

@@ -29,6 +29,7 @@ export class UiEvent {
   attachments?: { file: File; url: string }[];
   renderedContent?: any;
   a2uiData?: any;
+  textParts?: Array<{text: string, thought?: boolean}>;
   executableCode?: ExecutableCode;
   codeExecutionResult?: CodeExecutionResult;
   event!: Event;