hack: Subgraph streaming messages support (#394)

* hack: Subgraph streaming messages support

* cr

* cr

* cr

apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts

@@ -103,7 +103,8 @@ async function startProgrammerRun(input: {
       },
       ifNotExists: "create",
       streamResumable: true,
-      streamMode: ["values", "messages", "custom"],
+      streamSubgraphs: true,
+      streamMode: ["values", "messages", "custom", "events"],
     },
   );
 

apps/web/src/components/thread/messages/ai.tsx

@@ -123,14 +123,14 @@ function parseAnthropicStreamedToolCalls(
 
 export function mapToolMessageToActionStepProps(
   message: ToolMessage,
-  thread: { messages: Message[] },
+  threadMessages: Message[],
 ): ActionItemProps {
-  const toolCall: ToolCall | undefined = thread.messages
+  const toolCall: ToolCall | undefined = threadMessages
     .filter(isAIMessageSDK)
     .flatMap((m) => m.tool_calls ?? [])
     .find((tc) => tc.id === message.tool_call_id);
 
-  const aiMessage = thread.messages
+  const aiMessage = threadMessages
     .filter(isAIMessageSDK)
     .find((m) => m.tool_calls?.some((tc) => tc.id === message.tool_call_id));
   const reasoningText = aiMessage
@@ -234,12 +234,14 @@ export function AssistantMessage({
   handleRegenerate,
   forceRenderInterrupt = false,
   thread,
+  threadMessages,
 }: {
   message: Message | undefined;
   isLoading: boolean;
   handleRegenerate: (parentCheckpoint: Checkpoint | null | undefined) => void;
   forceRenderInterrupt?: boolean;
   thread: ReturnType<typeof useStream<Record<string, unknown>>>;
+  threadMessages: Message[];
 }) {
   const content = message?.content ?? [];
   const contentString = getContentString(content);
@@ -248,7 +250,7 @@ export function AssistantMessage({
     parseAsBoolean.withDefault(false),
   );
 
-  const messages = thread.messages;
+  const messages = threadMessages;
   const idx = message ? messages.findIndex((m) => m.id === message.id) : -1;
 
   const meta = message ? thread.getMessagesMetadata(message) : undefined;
@@ -409,7 +411,10 @@ export function AssistantMessage({
 
       if (correspondingToolResult) {
         // If we have a tool result, map it to action props
-        return mapToolMessageToActionStepProps(correspondingToolResult, thread);
+        return mapToolMessageToActionStepProps(
+          correspondingToolResult,
+          threadMessages,
+        );
       } else if (isRgTool) {
         const args = toolCall.args as RgToolArgs;
         return {
@@ -521,8 +526,8 @@ export function AssistantMessage({
   }
 
   const isLastMessage =
-    thread.messages[thread.messages.length - 1].id === message?.id;
+    threadMessages[threadMessages.length - 1].id === message?.id;
-  const hasNoAIOrToolMessages = !thread.messages.find(
+  const hasNoAIOrToolMessages = !threadMessages.find(
     (m) => m.type === "ai" || m.type === "tool",
   );
   const isToolResult = message?.type === "tool";

apps/web/src/components/v2/actions-renderer.tsx

@@ -14,6 +14,7 @@ import {
 import {
   DO_NOT_RENDER_ID_PREFIX,
   PLANNER_GRAPH_ID,
+  PROGRAMMER_GRAPH_ID,
 } from "@open-swe/shared/constants";
 import { Message } from "@langchain/langgraph-sdk";
 import { InitializeStep } from "../gen-ui/initialize-step";
@@ -23,6 +24,7 @@ import { GraphState, PlanItem } from "@open-swe/shared/open-swe/types";
 import { HumanResponse } from "@langchain/langgraph/prebuilt";
 import { LoadingActionsCardContent } from "./thread-view-loading";
 import { Interrupt } from "../thread/messages/interrupt";
+import { debounce } from "lodash";
 
 interface AcceptedPlanEventData {
   planTitle: string;
@@ -95,6 +97,38 @@ const getCustomNodeEventsFromMessages = (
     .flat();
 };
 
+function addMessagesToState(
+  existingMessages: Message[],
+  newMessages: Message[],
+): Message[] {
+  const existingIds = new Set(existingMessages.map((message) => message.id));
+  const uniqueNewMessages = newMessages.filter(
+    (message) => !message.id || !existingIds.has(message.id),
+  );
+  return [...existingMessages, ...uniqueNewMessages];
+}
+
+function isNodeEndMessagesUpdate(
+  data: unknown,
+): data is { output: { messages: Message[] } } {
+  return !!(
+    typeof data === "object" &&
+    data !== null &&
+    "output" in data &&
+    data.output &&
+    typeof data.output === "object" &&
+    "messages" in data.output &&
+    data.output.messages &&
+    Array.isArray(data.output.messages)
+  );
+}
+
+const REVIEWER_NODE_IDS = [
+  "take-review-actions",
+  "generate-review-actions",
+  "final-review",
+];
+
 export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
   graphId,
   threadId,
@@ -109,6 +143,8 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
   );
   const joinedRunId = useRef<string | undefined>(undefined);
   const [streamLoading, setStreamLoading] = useState(false);
+  const [mergedMessages, setMergedMessages] = useState<Message[]>([]);
+
   const stream = useStream<State>({
     apiUrl: process.env.NEXT_PUBLIC_API_URL,
     assistantId: graphId,
@@ -119,6 +155,18 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
         setCustomNodeEvents((prev) => [...prev, event]);
       }
     },
+    onLangChainEvent: (data) => {
+      if (
+        data.event === "on_chain_end" &&
+        data.metadata?.langgraph_node &&
+        REVIEWER_NODE_IDS.includes(data.metadata.langgraph_node as string) &&
+        data.data &&
+        isNodeEndMessagesUpdate(data.data)
+      ) {
+        const outputMessages = data.data.output.messages;
+        setMergedMessages((prev) => [...prev, ...outputMessages]);
+      }
+    },
     fetchStateHistory: false,
   });
 
@@ -147,7 +195,7 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
   }
 
   useEffect(() => {
-    const allCustomEvents = getCustomNodeEventsFromMessages(stream.messages);
+    const allCustomEvents = getCustomNodeEventsFromMessages(mergedMessages);
     if (!allCustomEvents?.length) {
       return;
     }
@@ -170,18 +218,18 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
 
       return prev;
     });
-  }, [stream.messages]);
+  }, [mergedMessages]);
 
   // Clear streamLoading as soon as we get any content (agent has started running)
   useEffect(() => {
     const hasContent =
-      (stream.messages && stream.messages.length > 0) ||
+      (mergedMessages && mergedMessages.length > 0) ||
       customNodeEvents.length > 0;
 
     if (hasContent && streamLoading) {
       setStreamLoading(false);
     }
-  }, [stream.messages, customNodeEvents, streamLoading]);
+  }, [mergedMessages, customNodeEvents, streamLoading]);
 
   // TODO: If the SDK changes go in, use this instead:
   // stream.joinStream(runId, undefined, { streamMode: ["values", "messages", "custom"]}).catch(console.error);
@@ -207,15 +255,15 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
   }, [onStreamReady, runId]); // Depend on runId instead of cancelRun to avoid infinite loops
 
   // Filter out human & do not render messages
-  const filteredMessages = stream.messages?.filter(
+  const filteredMessages = mergedMessages?.filter(
     (m) =>
       !isHumanMessageSDK(m) &&
       !(m.id && m.id.startsWith(DO_NOT_RENDER_ID_PREFIX)),
   );
   const isLastMessageHidden = !!(
-    stream.messages?.length > 0 &&
+    mergedMessages?.length > 0 &&
-    stream.messages[stream.messages.length - 1].id &&
+    mergedMessages[mergedMessages.length - 1].id &&
-    stream.messages[stream.messages.length - 1].id?.startsWith(
+    mergedMessages[mergedMessages.length - 1].id?.startsWith(
       DO_NOT_RENDER_ID_PREFIX,
     )
   );
@@ -242,6 +290,19 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
     }
   }, [stream.values, graphId]);
 
+  const debouncedSetMessages = useRef(
+    debounce((messages: Message[]) => {
+      setMergedMessages((prev) => addMessagesToState(prev, messages));
+    }, 100),
+  ).current;
+
+  useEffect(() => {
+    debouncedSetMessages(stream.messages);
+    return () => {
+      debouncedSetMessages.cancel();
+    };
+  }, [stream.messages, debouncedSetMessages]);
+
   if (streamLoading) {
     return <LoadingActionsCardContent />;
   }
@@ -260,6 +321,7 @@ export function ActionsRenderer<State extends PlannerGraphState | GraphState>({
         <AssistantMessage
           key={m.id}
           thread={stream as UseStream<Record<string, unknown>>}
+          threadMessages={mergedMessages}
           message={m}
           isLoading={false}
           handleRegenerate={() => {}}