From 6204bb984d15500268909201480b54515ac2536e Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Fri, 7 Mar 2025 21:26:48 +0100
Subject: [PATCH] Proposed fix

---
 agent/open-code/nodes/executor.ts     |  5 +++--
 agent/open-code/nodes/planner.ts      |  5 +++--
 src/components/thread/messages/ai.tsx | 19 +++++++++----------
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/agent/open-code/nodes/executor.ts b/agent/open-code/nodes/executor.ts
index 6276e1e..13f7eec 100644
--- a/agent/open-code/nodes/executor.ts
+++ b/agent/open-code/nodes/executor.ts
@@ -89,15 +89,16 @@ export async function executor(
     ],
   };
 
-  ui.write("proposed-change", {
+  const msg = ui.create("proposed-change", {
     toolCallId,
     change: updateFileContents,
     planItem,
   });
+  msg.additional_kwargs["message_id"] = aiMessage.id;
 
   return {
     messages: [aiMessage],
-    ui: ui.collect as OpenCodeUpdate["ui"],
+    ui: [msg],
     timestamp: Date.now(),
   };
 }
diff --git a/agent/open-code/nodes/planner.ts b/agent/open-code/nodes/planner.ts
index 2c5f94b..65e867c 100644
--- a/agent/open-code/nodes/planner.ts
+++ b/agent/open-code/nodes/planner.ts
@@ -55,11 +55,12 @@ export async function planner(
     ],
   };
 
-  ui.write("code-plan", {
+  const msg = ui.create("code-plan", {
     toolCallId,
     executedPlans,
     remainingPlans,
   });
+  msg.additional_kwargs["message_id"] = aiMessage.id;
 
   const toolMessage: ToolMessage = {
     type: "tool",
@@ -70,7 +71,7 @@ export async function planner(
 
   return {
     messages: [aiMessage, toolMessage],
-    ui: ui.collect as OpenCodeUpdate["ui"],
+    ui: [msg],
     timestamp: Date.now(),
   };
 }
diff --git a/src/components/thread/messages/ai.tsx b/src/components/thread/messages/ai.tsx
index 607d113..5784841 100644
--- a/src/components/thread/messages/ai.tsx
+++ b/src/components/thread/messages/ai.tsx
@@ -18,29 +18,28 @@ function CustomComponent({
   const [apiUrl] = useQueryParam("apiUrl", StringParam);
   const meta = thread.getMessagesMetadata(message);
   const seenState = meta?.firstSeenState;
-  const customComponent = seenState?.values.ui
+  const customComponents = seenState?.values.ui
     ?.slice()
-    .reverse()
-    .find(
+    .filter(
       ({ additional_kwargs }) =>
-        additional_kwargs.run_id === seenState.metadata?.run_id,
+        additional_kwargs.run_id === seenState.metadata?.run_id &&
+        (!additional_kwargs.message_id ||
+          additional_kwargs.message_id === message.id),
     );
 
-  if (!customComponent) {
-    return null;
-  }
-
+  if (!customComponents?.length) return null;
   return (
     <div key={message.id}>
-      {customComponent && (
+      {customComponents.map((customComponent) => (
         <LoadExternalComponent
+          key={customComponent.id}
           apiUrl={apiUrl ?? undefined}
           assistantId="agent"
           stream={thread}
           message={customComponent}
           meta={{ ui: customComponent }}
         />
-      )}
+      ))}
     </div>
   );
 }