diff --git a/agent/open-code/nodes/executor.ts b/agent/open-code/nodes/executor.ts
index 738c609..a8e59e4 100644
--- a/agent/open-code/nodes/executor.ts
+++ b/agent/open-code/nodes/executor.ts
@@ -101,15 +101,16 @@ export async function executor(
     ],
   };
 
-  ui.write("proposed-change", {
+  const msg = ui.create("proposed-change", {
     toolCallId,
     change: updateFileContents,
     planItem: nextPlanItem,
   });
+  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 6454483..5c9bf5d 100644
--- a/agent/open-code/nodes/planner.ts
+++ b/agent/open-code/nodes/planner.ts
@@ -70,11 +70,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",
@@ -85,7 +86,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>
   );
 }