From 9560b9814fe92e8a858907eac88aafaa992a6f20 Mon Sep 17 00:00:00 2001
From: doitundoit <taptypewrite@gmail.com>
Date: Sat, 3 Jan 2026 14:18:17 +0900
Subject: [PATCH] feat: improve Lit client robustness for message handling

- Add fallback mechanism to retrieve agent messages from history when status.message.parts is empty
- Handle string-encoded JSON data parts by parsing them
- Add debug logging (console.debug) for better traceability
---
 samples/client/lit/shell/client.ts | 38 +++++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/samples/client/lit/shell/client.ts b/samples/client/lit/shell/client.ts
index f0e94b8a..e81e017f 100644
--- a/samples/client/lit/shell/client.ts
+++ b/samples/client/lit/shell/client.ts
@@ -97,13 +97,45 @@ export class A2UIClient {
     }
 
     const result = (response as SendMessageSuccessResponse).result as Task;
-    if (result.kind === "task" && result.status.message?.parts) {
+    console.debug("Full Server Response Result:", JSON.stringify(result, null, 2));
+
+    let responseParts = result.status.message?.parts;
+
+    // Fallback: If no parts in status.message, check the last agent message in history
+    if (!responseParts && result.history && result.history.length > 0) {
+      // Iterate backwards to find the last agent message
+      for (let i = result.history.length - 1; i >= 0; i--) {
+        const msg = result.history[i];
+        if (msg.role === 'agent' && msg.parts && msg.parts.length > 0) {
+          responseParts = msg.parts;
+          console.debug("Found parts in history at index", i);
+          break;
+        }
+      }
+    }
+
+    if (result.kind === "task" && responseParts) {
       const messages: v0_8.Types.ServerToClientMessage[] = [];
-      for (const part of result.status.message.parts) {
+      for (const part of responseParts) {
+        console.debug("Client Received part:", JSON.stringify(part, null, 2));
+
         if (part.kind === 'data') {
-          messages.push(part.data as v0_8.Types.ServerToClientMessage);
+          let data = part.data;
+          // Handle string-encoded JSON data parts
+          if (typeof data === 'string') {
+            try {
+              data = JSON.parse(data);
+              console.debug("Parsed string data:", data);
+            } catch (e) {
+              console.error("Failed to parse part.data string:", e);
+            }
+          }
+          messages.push(data as v0_8.Types.ServerToClientMessage);
+        } else if (part.kind === 'text') {
+          console.debug("Ignored text part:", part.text);
         }
       }
+      console.debug("Final messages to process:", messages);
       return messages;
     }