From 24ce13936a6e86d0c6c4c50851e7d92b805ed18e Mon Sep 17 00:00:00 2001
From: Juan Tello <juan.tello@bigcommerce.com>
Date: Wed, 17 Dec 2025 17:38:53 -0500
Subject: [PATCH] feat: Enhance client message extraction from agent responses
 and update agent A2A client and request conversion imports.

---
 samples/agent/adk/orchestrator/agent.py       |  6 ++-
 .../agent/adk/orchestrator/agent_executor.py  |  1 +
 .../agent/adk/orchestrator/orchestrator.log   |  0
 samples/client/lit/shell/client.ts            | 45 +++++++++++++++++--
 4 files changed, 47 insertions(+), 5 deletions(-)
 create mode 100644 samples/agent/adk/orchestrator/orchestrator.log

diff --git a/samples/agent/adk/orchestrator/agent.py b/samples/agent/adk/orchestrator/agent.py
index ba878e08..60e8de8b 100644
--- a/samples/agent/adk/orchestrator/agent.py
+++ b/samples/agent/adk/orchestrator/agent.py
@@ -15,8 +15,12 @@
 import json
 import logging
 import os
-from typing import List
+from typing import Any, List
+from a2a.types import AgentCard
+from a2a.client.client import Consumer
 from a2a.client import A2ACardResolver
+from a2a.client.client import Client
+from a2a.client.middleware import ClientCallContext
 from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from google.adk.models.lite_llm import LiteLlm
 from google.adk.agents.llm_agent import LlmAgent
diff --git a/samples/agent/adk/orchestrator/agent_executor.py b/samples/agent/adk/orchestrator/agent_executor.py
index 366cbb0f..0286b326 100644
--- a/samples/agent/adk/orchestrator/agent_executor.py
+++ b/samples/agent/adk/orchestrator/agent_executor.py
@@ -30,6 +30,7 @@
     A2aAgentExecutorConfig,
     A2aAgentExecutor,
 )
+from google.adk.a2a.converters.request_converter import AgentRunRequest
 from a2a.types import AgentCapabilities, AgentCard, AgentExtension
 from a2ui.a2ui_extension import is_a2ui_part, try_activate_a2ui_extension, A2UI_EXTENSION_URI, STANDARD_CATALOG_ID, SUPPORTED_CATALOG_IDS_KEY, get_a2ui_agent_extension, A2UI_CLIENT_CAPABILITIES_KEY
 from google.adk.a2a.converters import event_converter
diff --git a/samples/agent/adk/orchestrator/orchestrator.log b/samples/agent/adk/orchestrator/orchestrator.log
new file mode 100644
index 00000000..e69de29b
diff --git a/samples/client/lit/shell/client.ts b/samples/client/lit/shell/client.ts
index f0e94b8a..88a9667b 100644
--- a/samples/client/lit/shell/client.ts
+++ b/samples/client/lit/shell/client.ts
@@ -97,16 +97,53 @@ export class A2UIClient {
     }
 
     const result = (response as SendMessageSuccessResponse).result as Task;
-    if (result.kind === "task" && result.status.message?.parts) {
-      const messages: v0_8.Types.ServerToClientMessage[] = [];
+    const messages: v0_8.Types.ServerToClientMessage[] = [];
+    
+    console.log("A2A Response:", JSON.stringify(result, null, 2));
+    
+    // Check for parts in the task status message (streaming/task-based responses)
+    if (result.kind === "task" && result.status?.message?.parts) {
+      console.log("Found parts in task status message");
       for (const part of result.status.message.parts) {
         if (part.kind === 'data') {
           messages.push(part.data as v0_8.Types.ServerToClientMessage);
         }
       }
-      return messages;
+      if (messages.length > 0) return messages;
+    }
+    
+    // Check for parts in artifacts (blocking responses)
+    if (result.artifacts && result.artifacts.length > 0) {
+      console.log("Found artifacts:", result.artifacts.length);
+      for (const artifact of result.artifacts) {
+        if (artifact.parts) {
+          for (const part of artifact.parts) {
+            console.log("Artifact part kind:", part.kind);
+            if (part.kind === 'data') {
+              messages.push(part.data as v0_8.Types.ServerToClientMessage);
+            }
+          }
+        }
+      }
+      if (messages.length > 0) return messages;
+    }
+    
+    // Check for parts in the last history message (agent response)
+    if (result.history && result.history.length > 0) {
+      console.log("Found history:", result.history.length);
+      const lastMessage = result.history[result.history.length - 1];
+      console.log("Last message role:", lastMessage.role);
+      if (lastMessage.role === 'agent' && lastMessage.parts) {
+        for (const part of lastMessage.parts) {
+          console.log("History part kind:", part.kind);
+          if (part.kind === 'data') {
+            messages.push(part.data as v0_8.Types.ServerToClientMessage);
+          }
+        }
+      }
     }
 
-    return [];
+    console.log("Total messages found:", messages.length);
+    return messages;
   }
 }