From caf222235fd6892ff0da1b9b22771a7d7bdf6dcf Mon Sep 17 00:00:00 2001 From: Brace Sproul Date: Fri, 11 Jul 2025 13:58:10 -0700 Subject: [PATCH] hack: Subgraph streaming messages support (#394) * hack: Subgraph streaming messages support * cr * cr * cr --- .../src/graphs/planner/nodes/proposed-plan.ts | 3 +- .../web/src/components/thread/messages/ai.tsx | 19 +++-- .../src/components/v2/actions-renderer.tsx | 78 +++++++++++++++++-- 3 files changed, 84 insertions(+), 16 deletions(-) diff --git a/apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts b/apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts index 291b41e8..ef49d722 100644 --- a/apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts +++ b/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"], }, ); diff --git a/apps/web/src/components/thread/messages/ai.tsx b/apps/web/src/components/thread/messages/ai.tsx index 992ee7e2..675ed570 100644 --- a/apps/web/src/components/thread/messages/ai.tsx +++ b/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>>; + 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; - const hasNoAIOrToolMessages = !thread.messages.find( + threadMessages[threadMessages.length - 1].id === message?.id; + const hasNoAIOrToolMessages = !threadMessages.find( (m) => m.type === "ai" || m.type === "tool", ); const isToolResult = message?.type === "tool"; diff --git a/apps/web/src/components/v2/actions-renderer.tsx b/apps/web/src/components/v2/actions-renderer.tsx index 675ad094..2eec3466 100644 --- a/apps/web/src/components/v2/actions-renderer.tsx +++ b/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({ graphId, threadId, @@ -109,6 +143,8 @@ export function ActionsRenderer({ ); const joinedRunId = useRef(undefined); const [streamLoading, setStreamLoading] = useState(false); + const [mergedMessages, setMergedMessages] = useState([]); + const stream = useStream({ apiUrl: process.env.NEXT_PUBLIC_API_URL, assistantId: graphId, @@ -119,6 +155,18 @@ export function ActionsRenderer({ 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({ } useEffect(() => { - const allCustomEvents = getCustomNodeEventsFromMessages(stream.messages); + const allCustomEvents = getCustomNodeEventsFromMessages(mergedMessages); if (!allCustomEvents?.length) { return; } @@ -170,18 +218,18 @@ export function ActionsRenderer({ 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({ }, [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 && - stream.messages[stream.messages.length - 1].id && - stream.messages[stream.messages.length - 1].id?.startsWith( + mergedMessages?.length > 0 && + mergedMessages[mergedMessages.length - 1].id && + mergedMessages[mergedMessages.length - 1].id?.startsWith( DO_NOT_RENDER_ID_PREFIX, ) ); @@ -242,6 +290,19 @@ export function ActionsRenderer({ } }, [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 ; } @@ -260,6 +321,7 @@ export function ActionsRenderer({ >} + threadMessages={mergedMessages} message={m} isLoading={false} handleRegenerate={() => {}}