From c82bd4d6ab18ac04db9752e4bc7378f5b3027bc9 Mon Sep 17 00:00:00 2001
From: Yifan Wang <wanyif@google.com>
Date: Tue, 2 Jun 2026 10:12:56 -0700
Subject: [PATCH] Support rendering A2UI json directly from content

---
 .../components/chat/chat.component.spec.ts    | 44 +++++++++++++++++++
 src/app/components/chat/chat.component.ts     | 44 +++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/src/app/components/chat/chat.component.spec.ts b/src/app/components/chat/chat.component.spec.ts
index 34d48920..650da130 100644
--- a/src/app/components/chat/chat.component.spec.ts
+++ b/src/app/components/chat/chat.component.spec.ts
@@ -1383,6 +1383,50 @@ describe('ChatComponent', () => {
         });
   });
 
+  describe('extractA2uiJsonFromText', () => {
+    it('should do nothing if message has no text', () => {
+      const uiEvent = new UiEvent({role: 'bot', event: {} as any});
+      (component as any).extractA2uiJsonFromText(uiEvent);
+      expect(uiEvent.a2uiData).toBeUndefined();
+    });
+
+    it('should do nothing if text has no <a2ui-json> tags', () => {
+      const uiEvent = new UiEvent({
+        role: 'bot',
+        text: 'hello world',
+        event: {} as any
+      });
+      (component as any).extractA2uiJsonFromText(uiEvent);
+      expect(uiEvent.a2uiData).toBeUndefined();
+      expect(uiEvent.text).toBe('hello world');
+    });
+
+    it('should extract and parse inline <a2ui-json> block, and strip it from text', () => {
+      const payload = [{beginRendering: {surfaceId: 'cloud_dash'}}];
+      const uiEvent = new UiEvent({
+        role: 'bot',
+        text: `Here is the UI:\n<a2ui-json>\n${JSON.stringify(payload)}\n</a2ui-json>\nEnjoy!`,
+        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!');
+    });
+
+    it('should keep tags and log warning if JSON parsing fails', () => {
+      const uiEvent = new UiEvent({
+        role: 'bot',
+        text: `broken tags: <a2ui-json>{broken-json}</a2ui-json>`,
+        event: {} as any
+      });
+
+      (component as any).extractA2uiJsonFromText(uiEvent);
+      expect(uiEvent.a2uiData).toBeUndefined();
+      expect(uiEvent.text).toBe('broken tags: <a2ui-json>{broken-json}</a2ui-json>');
+    });
+  });
+
   describe('refreshLatestSession', () => {
     beforeEach(() => {
       mockAgentService.listAppsResponse.next([TEST_APP_1_NAME]);
diff --git a/src/app/components/chat/chat.component.ts b/src/app/components/chat/chat.component.ts
index fd5f7fc2..a238cd61 100644
--- a/src/app/components/chat/chat.component.ts
+++ b/src/app/components/chat/chat.component.ts
@@ -1541,6 +1541,48 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
     return a2uiData;
   }
 
+  private extractA2uiJsonFromText(uiEvent: UiEvent) {
+    if (!uiEvent.text) return;
+
+    const startTag = '<a2ui-json>';
+    const endTag = '</a2ui-json>';
+
+    const startIndex = uiEvent.text.indexOf(startTag);
+    if (startIndex === -1) return;
+
+    const endIndex = uiEvent.text.indexOf(endTag, startIndex + startTag.length);
+    if (endIndex === -1) return;
+
+    const jsonStr = uiEvent.text.substring(startIndex + startTag.length, endIndex).trim();
+    try {
+      let parsed = JSON.parse(jsonStr);
+      if (!Array.isArray(parsed)) {
+        parsed = [parsed];
+      }
+
+      const a2uiData: any = {};
+      parsed.forEach((msg: any) => {
+        if (msg.beginRendering) {
+          a2uiData.beginRendering = msg;
+        } else if (msg.surfaceUpdate) {
+          a2uiData.surfaceUpdate = msg;
+        } else if (msg.dataModelUpdate) {
+          a2uiData.dataModelUpdate = msg;
+        }
+      });
+
+      uiEvent.a2uiData = a2uiData;
+
+      // Strip the <a2ui-json> block from uiEvent.text
+      const beforeText = uiEvent.text.substring(0, startIndex);
+      const afterText = uiEvent.text.substring(endIndex + endTag.length);
+      uiEvent.text = (beforeText + afterText).trim();
+
+    } catch (e) {
+      console.warn('Failed to parse inline <a2ui-json> block from text:', e);
+    }
+  }
+
   private updateRedirectUri(urlString: string, newRedirectUri: string): string {
     try {
       const url = new URL(urlString);
@@ -2153,6 +2195,8 @@ export class ChatComponent implements OnInit, AfterViewInit, OnDestroy {
       this.processPartIntoMessage(part, event, uiEvent);
     });
 
+    this.extractA2uiJsonFromText(uiEvent);
+
     return uiEvent;
   }