From 23418cab1cf161901907adc4070966a57f109c0c Mon Sep 17 00:00:00 2001 From: Brace Sproul Date: Thu, 19 Jun 2025 11:53:56 -0700 Subject: [PATCH] feat: Custom events (#186) * feat: Custom events * implementing initialization component synced with initialize node * handle irrelevant/skipped events * fix positioning, add collapse, rm comments * lint format * fix tsconfig, refactor DRY eventEmit, collapse initialize-step * auto collapse init-step after completion * revert yarn lock to main * formatting * cr * enable persisting init steps after stream * format * cr * cr * cr --------- Co-authored-by: starmorph --- .../planner/nodes/generate-message/index.ts | 3 +- .../src/graphs/planner/nodes/generate-plan.ts | 3 +- .../src/graphs/planner/nodes/prepare-state.ts | 5 +- .../src/graphs/shared/initialize-sandbox.ts | 310 ++++++++++++++++-- apps/open-swe/src/utils/langgraph-client.ts | 4 +- .../src/utils/message/filter-hidden.ts | 5 + .../src/components/gen-ui/initialize-step.tsx | 144 ++++---- .../hooks/useInterruptedActions.tsx | 1 + .../agent-inbox/hooks/useProposedPlan.tsx | 2 + apps/web/src/components/thread/index.tsx | 32 +- .../src/components/v2/actions-renderer.tsx | 86 ++++- apps/web/src/components/v2/terminal-input.tsx | 55 ++-- apps/web/src/components/v2/thread-view.tsx | 22 +- apps/web/src/lib/ensure-tool-responses.ts | 3 +- apps/web/src/providers/Stream.tsx | 21 +- packages/shared/src/constants.ts | 2 + .../shared/src/open-swe/custom-node-events.ts | 74 +++++ 17 files changed, 621 insertions(+), 151 deletions(-) create mode 100644 apps/open-swe/src/utils/message/filter-hidden.ts create mode 100644 packages/shared/src/open-swe/custom-node-events.ts diff --git a/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts b/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts index 9fda477f..c1adf9dd 100644 --- a/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts +++ b/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts @@ -14,6 +14,7 @@ import { import { SYSTEM_PROMPT } from "./prompt.js"; import { getRepoAbsolutePath } from "@open-swe/shared/git"; import { getMissingMessages } from "../../../../utils/github/issue-messages.js"; +import { filterHiddenMessages } from "../../../../utils/message/filter-hidden.js"; import { getTaskPlanFromIssue } from "../../../../utils/github/issue-task.js"; const logger = createLogger(LogLevel.INFO, "GeneratePlanningMessageNode"); @@ -62,7 +63,7 @@ export async function generateAction( taskPlan: latestTaskPlan ?? state.taskPlan, }), }, - ...state.messages, + ...filterHiddenMessages(state.messages), ...missingMessages, ]); diff --git a/apps/open-swe/src/graphs/planner/nodes/generate-plan.ts b/apps/open-swe/src/graphs/planner/nodes/generate-plan.ts index 634014ad..1d50de2d 100644 --- a/apps/open-swe/src/graphs/planner/nodes/generate-plan.ts +++ b/apps/open-swe/src/graphs/planner/nodes/generate-plan.ts @@ -12,6 +12,7 @@ import { isFollowupRequest, } from "../utils/followup.js"; import { stopSandbox } from "../../../utils/sandbox.js"; +import { filterHiddenMessages } from "../../../utils/message/filter-hidden.js"; const systemPrompt = `You are operating as a terminal-based agentic coding assistant built by LangChain. It wraps LLM models to enable natural language interaction with a local codebase. You are expected to be precise, safe, and helpful. {FOLLOWUP_MESSAGE_PROMPT} @@ -80,7 +81,7 @@ export async function generatePlan( role: "system", content: formatSystemPrompt(state), }, - ...state.messages, + ...filterHiddenMessages(state.messages), ...(optionalToolMessage ? [optionalToolMessage] : []), ]); diff --git a/apps/open-swe/src/graphs/planner/nodes/prepare-state.ts b/apps/open-swe/src/graphs/planner/nodes/prepare-state.ts index fc0df7aa..406867b2 100644 --- a/apps/open-swe/src/graphs/planner/nodes/prepare-state.ts +++ b/apps/open-swe/src/graphs/planner/nodes/prepare-state.ts @@ -18,6 +18,7 @@ import { getMessageContentFromIssue, getUntrackedComments, } from "../../../utils/github/issue-messages.js"; +import { filterHiddenMessages } from "../../../utils/message/filter-hidden.js"; export async function prepareGraphState( state: PlannerGraphState, @@ -92,8 +93,8 @@ export async function prepareGraphState( }); } - // Remove all messages not marked as summaryMessage, and not human messages. - const removedNonSummaryMessages = state.messages + // Remove all messages not marked as summaryMessage, hidden, and not human messages. + const removedNonSummaryMessages = filterHiddenMessages(state.messages) .filter((m) => !m.additional_kwargs?.summaryMessage && !isHumanMessage(m)) .map((m: BaseMessage) => new RemoveMessage({ id: m.id ?? "" })); const summaryMessage = new AIMessage({ diff --git a/apps/open-swe/src/graphs/shared/initialize-sandbox.ts b/apps/open-swe/src/graphs/shared/initialize-sandbox.ts index 309676b1..7e234b44 100644 --- a/apps/open-swe/src/graphs/shared/initialize-sandbox.ts +++ b/apps/open-swe/src/graphs/shared/initialize-sandbox.ts @@ -1,3 +1,4 @@ +import { v4 as uuidv4 } from "uuid"; import { getRepoAbsolutePath } from "@open-swe/shared/git"; import { getGitHubTokensFromConfig } from "../../utils/github-tokens.js"; import { GraphConfig, TargetRepository } from "@open-swe/shared/open-swe/types"; @@ -10,7 +11,16 @@ import { pullLatestChanges, } from "../../utils/github/git.js"; import { getCodebaseTree } from "../../utils/tree.js"; -import { SNAPSHOT_NAME } from "@open-swe/shared/constants"; +import { + DO_NOT_RENDER_ID_PREFIX, + SNAPSHOT_NAME, +} from "@open-swe/shared/constants"; +import { + CustomNodeEvent, + INITIALIZE_NODE_ID, +} from "@open-swe/shared/open-swe/custom-node-events"; +import { Sandbox } from "@daytonaio/sdk"; +import { AIMessage, BaseMessage } from "@langchain/core/messages"; const logger = createLogger(LogLevel.INFO, "InitializeSandbox"); @@ -19,6 +29,7 @@ type InitializeSandboxState = { branchName: string; sandboxSessionId?: string; codebaseTree?: string; + messages?: BaseMessage[]; }; export async function initializeSandbox( @@ -28,65 +39,302 @@ export async function initializeSandbox( const { githubInstallationToken } = getGitHubTokensFromConfig(config); const { sandboxSessionId, targetRepository, branchName } = state; const absoluteRepoDir = getRepoAbsolutePath(targetRepository); + const repoName = `${targetRepository.owner}/${targetRepository.repo}`; + + const events: CustomNodeEvent[] = []; + const emitStepEvent = ( + base: CustomNodeEvent, + status: "pending" | "success" | "error" | "skipped", + error?: string, + ) => { + const event = { + ...base, + createdAt: new Date().toISOString(), + data: { + ...base.data, + status, + ...(error ? { error } : {}), + }, + }; + events.push(event); + try { + config.writer?.(event); + } catch (err) { + logger.error("Failed to emit custom event", { event, err }); + } + }; + const createEventsMessage = () => [ + new AIMessage({ + id: `${DO_NOT_RENDER_ID_PREFIX}${uuidv4()}`, + content: "", + additional_kwargs: { + hidden: true, + customNodeEvents: events, + }, + }), + ]; + + if (!sandboxSessionId) { + emitStepEvent( + { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: uuidv4(), + action: "Resuming Sandbox", + data: { + status: "skipped", + branch: branchName, + repo: repoName, + }, + }, + "skipped", + ); + emitStepEvent( + { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: uuidv4(), + action: "Pulling latest changes", + data: { + status: "skipped", + branch: branchName, + repo: repoName, + }, + }, + "skipped", + ); + } if (sandboxSessionId) { - try { - logger.info("Sandbox session ID exists. Resuming", { + const resumeSandboxActionId = uuidv4(); + const baseResumeSandboxAction: CustomNodeEvent = { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: resumeSandboxActionId, + action: "Resuming Sandbox", + data: { + status: "pending", sandboxSessionId, - }); - // Resume the sandbox if the session ID is in the config. + branch: branchName, + repo: repoName, + }, + }; + emitStepEvent(baseResumeSandboxAction, "pending"); + + try { const existingSandbox = await daytonaClient().get(sandboxSessionId); - await pullLatestChanges(absoluteRepoDir, existingSandbox); - const codebaseTree = await getCodebaseTree(existingSandbox.id); - return { - sandboxSessionId: existingSandbox.id, - codebaseTree, + emitStepEvent(baseResumeSandboxAction, "success"); + + const pullLatestChangesActionId = uuidv4(); + const basePullLatestChangesAction: CustomNodeEvent = { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: pullLatestChangesActionId, + action: "Pulling latest changes", + data: { + status: "pending", + sandboxSessionId, + branch: branchName, + repo: repoName, + }, }; - } catch (e) { - // Error thrown, log it and continue. Will create a new sandbox session since the resumption failed. - logger.error("Failed to get sandbox session", e); + emitStepEvent(basePullLatestChangesAction, "pending"); + + const pullChangesRes = await pullLatestChanges( + absoluteRepoDir, + existingSandbox, + ); + if (!pullChangesRes || pullChangesRes.exitCode !== 0) { + emitStepEvent( + basePullLatestChangesAction, + "error", + "Failed to pull latest changes. Please check your repository connection.", + ); + throw new Error("Failed to pull latest changes."); + } + emitStepEvent(basePullLatestChangesAction, "success"); + + const generateCodebaseTreeActionId = uuidv4(); + const baseGenerateCodebaseTreeAction: CustomNodeEvent = { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: generateCodebaseTreeActionId, + action: "Generating codebase tree", + data: { + status: "pending", + sandboxSessionId, + branch: branchName, + repo: repoName, + }, + }; + emitStepEvent(baseGenerateCodebaseTreeAction, "pending"); + try { + const codebaseTree = await getCodebaseTree(existingSandbox.id); + emitStepEvent(baseGenerateCodebaseTreeAction, "success"); + return { + sandboxSessionId: existingSandbox.id, + codebaseTree, + messages: createEventsMessage(), + }; + } catch { + emitStepEvent( + baseGenerateCodebaseTreeAction, + "error", + "Failed to generate codebase tree. Please try again later.", + ); + throw new Error("Failed to generate codebase tree."); + } + } catch { + emitStepEvent( + baseResumeSandboxAction, + "error", + "Failed to resume sandbox. A new environment will be created.", + ); } } - logger.info("Creating sandbox..."); - const sandbox = await daytonaClient().create({ - image: SNAPSHOT_NAME, - }); + // Creating Sandbox + const createSandboxActionId = uuidv4(); + const baseCreateSandboxAction: CustomNodeEvent = { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: createSandboxActionId, + action: "Creating Sandbox", + data: { + status: "pending", + sandboxSessionId: null, + branch: branchName, + repo: repoName, + }, + }; - const res = await cloneRepo(sandbox, targetRepository, { - githubInstallationToken, - stateBranchName: branchName, - }); - if (res.exitCode !== 0) { - throw new Error(`Failed to clone repository.\n${res.result}`); + emitStepEvent(baseCreateSandboxAction, "pending"); + let sandbox: Sandbox; + try { + sandbox = await daytonaClient().create({ image: SNAPSHOT_NAME }); + emitStepEvent(baseCreateSandboxAction, "success"); + } catch { + emitStepEvent( + baseCreateSandboxAction, + "error", + "Failed to create sandbox environment. Please try again later.", + ); + throw new Error("Failed to create sandbox environment."); } - logger.info("Repository cloned successfully."); + // Cloning repository + const cloneRepoActionId = uuidv4(); + const baseCloneRepoAction: CustomNodeEvent = { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: cloneRepoActionId, + action: "Cloning repository", + data: { + status: "pending", + sandboxSessionId: sandbox.id, + branch: branchName, + repo: repoName, + }, + }; + emitStepEvent(baseCloneRepoAction, "pending"); + const cloneRepoRes = await cloneRepo(sandbox, targetRepository, { + githubInstallationToken, + stateBranchName: state.branchName, + }); + if (cloneRepoRes.exitCode !== 0) { + emitStepEvent( + baseCloneRepoAction, + "error", + "Failed to clone repository. Please check your repo URL and permissions.", + ); + throw new Error("Failed to clone repository."); + } + emitStepEvent(baseCloneRepoAction, "success"); + + // Configuring git user + const configureGitUserActionId = uuidv4(); + const baseConfigureGitUserAction: CustomNodeEvent = { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: configureGitUserActionId, + action: "Configuring git user", + data: { + status: "pending", + sandboxSessionId: sandbox.id, + branch: branchName, + repo: repoName, + }, + }; + emitStepEvent(baseConfigureGitUserAction, "pending"); - logger.info(`Configuring git user for repository at "${absoluteRepoDir}"...`); await configureGitUserInRepo(absoluteRepoDir, sandbox, { githubInstallationToken, owner: targetRepository.owner, repo: targetRepository.repo, }); - logger.info("Git user configured successfully."); + emitStepEvent(baseConfigureGitUserAction, "success"); + // Checking out branch + const checkoutBranchActionId = uuidv4(); + const baseCheckoutBranchAction: CustomNodeEvent = { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: checkoutBranchActionId, + action: "Checking out branch", + data: { + status: "pending", + sandboxSessionId: sandbox.id, + branch: branchName, + repo: repoName, + }, + }; + emitStepEvent(baseCheckoutBranchAction, "pending"); const checkoutBranchRes = await checkoutBranch( absoluteRepoDir, branchName, sandbox, ); - if (!checkoutBranchRes) { - // TODO: This should probably be an interrupt. - logger.error("Failed to checkout branch."); - throw new Error("Failed to checkout branch"); + emitStepEvent( + baseCheckoutBranchAction, + "error", + "Failed to checkout branch. Please check your branch name.", + ); + throw new Error("Failed to checkout branch."); } + emitStepEvent(baseCheckoutBranchAction, "success"); - const codebaseTree = await getCodebaseTree(sandbox.id); + // Generating codebase tree + const generateCodebaseTreeActionId = uuidv4(); + const baseGenerateCodebaseTreeAction: CustomNodeEvent = { + nodeId: INITIALIZE_NODE_ID, + createdAt: new Date().toISOString(), + actionId: generateCodebaseTreeActionId, + action: "Generating codebase tree", + data: { + status: "pending", + sandboxSessionId: sandbox.id, + branch: branchName, + repo: repoName, + }, + }; + emitStepEvent(baseGenerateCodebaseTreeAction, "pending"); + let codebaseTree = undefined; + try { + codebaseTree = await getCodebaseTree(sandbox.id); + emitStepEvent(baseGenerateCodebaseTreeAction, "success"); + } catch (_) { + emitStepEvent( + baseGenerateCodebaseTreeAction, + "error", + "Failed to generate codebase tree.", + ); + } return { sandboxSessionId: sandbox.id, + targetRepository, codebaseTree, + messages: createEventsMessage(), }; } diff --git a/apps/open-swe/src/utils/langgraph-client.ts b/apps/open-swe/src/utils/langgraph-client.ts index 4c91a377..97f9c9e6 100644 --- a/apps/open-swe/src/utils/langgraph-client.ts +++ b/apps/open-swe/src/utils/langgraph-client.ts @@ -4,6 +4,8 @@ export function createLangGraphClient(options?: { defaultHeaders?: Record; includeApiKey?: boolean; }) { + // TODO: Remove the need for this after issues with port are resolved. + const productionUrl = process.env.LANGGRAPH_PROD_URL; const port = process.env.PORT ?? "2024"; if (options?.includeApiKey && !process.env.LANGGRAPH_API_KEY) { throw new Error("LANGGRAPH_API_KEY not found"); @@ -12,7 +14,7 @@ export function createLangGraphClient(options?: { ...(options?.includeApiKey && { apiKey: process.env.LANGGRAPH_API_KEY, }), - apiUrl: `http://localhost:${port}`, + apiUrl: productionUrl ?? `http://localhost:${port}`, defaultHeaders: options?.defaultHeaders, }); } diff --git a/apps/open-swe/src/utils/message/filter-hidden.ts b/apps/open-swe/src/utils/message/filter-hidden.ts new file mode 100644 index 00000000..9e402f07 --- /dev/null +++ b/apps/open-swe/src/utils/message/filter-hidden.ts @@ -0,0 +1,5 @@ +import { BaseMessage } from "@langchain/core/messages"; + +export function filterHiddenMessages(messages: BaseMessage[]): BaseMessage[] { + return messages.filter((message) => !message.additional_kwargs?.hidden); +} diff --git a/apps/web/src/components/gen-ui/initialize-step.tsx b/apps/web/src/components/gen-ui/initialize-step.tsx index eed5448a..b895dfc8 100644 --- a/apps/web/src/components/gen-ui/initialize-step.tsx +++ b/apps/web/src/components/gen-ui/initialize-step.tsx @@ -1,6 +1,4 @@ "use client"; - -import "../app/globals.css"; import { Loader2, CheckCircle, @@ -8,36 +6,43 @@ import { GitBranch, MessageSquare, FileText, + ChevronDown, } from "lucide-react"; -import { useState } from "react"; - -type Step = { - name: string; - status: "waiting" | "generating" | "success" | "error"; - error?: string; -}; +import { useState, useEffect, useRef } from "react"; +import { cn } from "@/lib/utils"; +import { Step } from "@open-swe/shared/open-swe/custom-node-events"; +import { Button } from "../ui/button"; type InitializeStepProps = { status: "loading" | "generating" | "done"; success?: boolean; steps?: Step[]; - reasoningText?: string; - summaryText?: string; + collapse?: boolean; }; export function InitializeStep({ status, success, steps, - reasoningText, - summaryText, + collapse: collapseProp, }: InitializeStepProps) { - const [showReasoning, setShowReasoning] = useState(false); - const [showSummary, setShowSummary] = useState(false); + const [collapsed, setCollapsed] = useState(collapseProp ?? false); + const wasDone = useRef(false); + + // Auto-collapse when status is 'done' and success is true + useEffect(() => { + if (status === "done" && success && !collapsed && !wasDone.current) { + setCollapsed(true); + wasDone.current = true; + } + if (status !== "done") { + wasDone.current = false; + } + }, [status, success, collapsed]); const stepStatusIcon = { waiting: ( -
+
), generating: , success: , @@ -48,7 +53,9 @@ export function InitializeStep({ switch (status) { case "loading": return ( -
+
); case "generating": return ; @@ -74,72 +81,69 @@ export function InitializeStep({ return (
- {reasoningText && ( -
- - {showReasoning && ( -

- {reasoningText} -

- )} -
- )} - -
+ {/* Collapse/Expand Icon */} +
{getStatusText()} {getStatusIcon()} +
- - {steps && (status === "generating" || status === "done") && ( + {/* Only render the rest if not collapsed */} + {!collapsed && steps && steps.length > 0 && (
    - {steps.map((step, index) => ( -
  • - {stepStatusIcon[step.status]} - step.status !== "skipped") + .map((step, index) => ( +
  • - {step.name} - - {step.error && ( - - ({step.error}) + + {stepStatusIcon[ + step.status as keyof typeof stepStatusIcon + ] ?? ( +
    + )} - )} -
    • - ))} + + {step.name} + + {step.error && ( + + ({step.error}) + + )} + + ))}
)} - - {summaryText && status === "done" && ( -
- - {showSummary && ( -

- {summaryText} -

- )} -
- )}
); } diff --git a/apps/web/src/components/thread/agent-inbox/hooks/useInterruptedActions.tsx b/apps/web/src/components/thread/agent-inbox/hooks/useInterruptedActions.tsx index ecdebbab..7ec66a79 100644 --- a/apps/web/src/components/thread/agent-inbox/hooks/useInterruptedActions.tsx +++ b/apps/web/src/components/thread/agent-inbox/hooks/useInterruptedActions.tsx @@ -94,6 +94,7 @@ export default function useInterruptedActions({ config: { recursion_limit: 400, }, + streamResumable: true, }, ); diff --git a/apps/web/src/components/thread/agent-inbox/hooks/useProposedPlan.tsx b/apps/web/src/components/thread/agent-inbox/hooks/useProposedPlan.tsx index 210fe844..d20acb85 100644 --- a/apps/web/src/components/thread/agent-inbox/hooks/useProposedPlan.tsx +++ b/apps/web/src/components/thread/agent-inbox/hooks/useProposedPlan.tsx @@ -49,6 +49,7 @@ export function useProposedPlan( config: { recursion_limit: 400, }, + streamResumable: true, }, ); }; @@ -69,6 +70,7 @@ export function useProposedPlan( config: { recursion_limit: 400, }, + streamResumable: true, }, ); }; diff --git a/apps/web/src/components/thread/index.tsx b/apps/web/src/components/thread/index.tsx index b324ff91..2c730f47 100644 --- a/apps/web/src/components/thread/index.tsx +++ b/apps/web/src/components/thread/index.tsx @@ -8,10 +8,7 @@ import { Button } from "../ui/button"; import { Checkpoint, Message } from "@langchain/langgraph-sdk"; import { AssistantMessage, AssistantMessageLoading } from "./messages/ai"; import { HumanMessage } from "./messages/human"; -import { - DO_NOT_RENDER_ID_PREFIX, - ensureToolCallsHaveResponses, -} from "@/lib/ensure-tool-responses"; +import { ensureToolCallsHaveResponses } from "@/lib/ensure-tool-responses"; import { LangGraphLogoSVG } from "../icons/langgraph"; import { TooltipIconButton } from "../ui/tooltip-icon-button"; import { @@ -52,6 +49,11 @@ import { TaskPlanView } from "../tasks"; import { useTaskPlan } from "../tasks/useTaskPlan"; import { isProposedPlanInterrupt } from "@/lib/plan-utils"; import { HumanResponse } from "@langchain/langgraph/prebuilt"; +import { + INITIALIZE_NODE_ID, + mapCustomEventsToSteps, +} from "@open-swe/shared/open-swe/custom-node-events"; +import { DO_NOT_RENDER_ID_PREFIX } from "@open-swe/shared/constants"; function StickyToBottomContent(props: { content: ReactNode; @@ -151,6 +153,7 @@ export function Thread() { const stream = useStreamContext(); const messages = stream.messages; const isLoading = stream.isLoading; + const customEvents = stream.customEvents; const lastError = useRef(undefined); @@ -321,6 +324,19 @@ export function Thread() { ); const isLastMessageHuman = messages[messages.length - 1]?.type === "human"; + const initializeEvents = customEvents.filter( + (e) => e.nodeId === INITIALIZE_NODE_ID, + ); + + const steps = mapCustomEventsToSteps(initializeEvents); + const allSuccess = + steps.length > 0 && steps.every((s) => s.status === "success"); + + let initStatus: "loading" | "generating" | "done" = "generating"; + if (allSuccess) { + initStatus = "done"; + } + return (
@@ -513,9 +529,11 @@ export function Thread() { thread={stream} /> )} - {isLoading && !firstTokenReceived && ( - - )} + {isLoading && + !firstTokenReceived && + initializeEvents.length === 0 && ( + + )} } footer={ diff --git a/apps/web/src/components/v2/actions-renderer.tsx b/apps/web/src/components/v2/actions-renderer.tsx index dbc402d5..336da898 100644 --- a/apps/web/src/components/v2/actions-renderer.tsx +++ b/apps/web/src/components/v2/actions-renderer.tsx @@ -1,8 +1,17 @@ -import { isHumanMessageSDK } from "@/lib/langchain-messages"; +import { isAIMessageSDK, isHumanMessageSDK } from "@/lib/langchain-messages"; import { UseStream, useStream } from "@langchain/langgraph-sdk/react"; import { AssistantMessage } from "../thread/messages/ai"; -import { Dispatch, SetStateAction, useEffect, useRef } from "react"; +import { Dispatch, SetStateAction, useEffect, useRef, useState } from "react"; import { ManagerGraphState } from "@open-swe/shared/open-swe/manager/types"; +import { + isCustomNodeEvent, + CustomNodeEvent, + INITIALIZE_NODE_ID, + mapCustomEventsToSteps, +} from "@open-swe/shared/open-swe/custom-node-events"; +import { DO_NOT_RENDER_ID_PREFIX } from "@open-swe/shared/constants"; +import { Message } from "@langchain/langgraph-sdk"; +import { InitializeStep } from "../gen-ui/initialize-step"; import { PlannerGraphState } from "@open-swe/shared/open-swe/planner/types"; import { GraphState } from "@open-swe/shared/open-swe/types"; @@ -17,6 +26,28 @@ interface ActionsRendererProps { setSelectedTab?: Dispatch>; } +const getCustomNodeEventsFromMessages = ( + messages?: Message[], + nodeId?: string, +): CustomNodeEvent[] => { + if (!messages) return []; + + return messages + .filter((m) => isAIMessageSDK(m)) + .filter((m) => { + const events = m.additional_kwargs?.customNodeEvents as + | CustomNodeEvent[] + | undefined; + if (!events?.length) return false; + if (!nodeId) { + return true; + } + return events.some((e) => e.nodeId === nodeId); + }) + .map((m) => m.additional_kwargs?.customNodeEvents as CustomNodeEvent[]) + .flat(); +}; + export function ActionsRenderer({ graphId, threadId, @@ -25,23 +56,60 @@ export function ActionsRenderer({ programmerSession, setSelectedTab, }: ActionsRendererProps) { + const [customNodeEvents, setCustomNodeEvents] = useState( + [], + ); const stream = useStream({ apiUrl: process.env.NEXT_PUBLIC_API_URL, assistantId: graphId, reconnectOnMount: true, threadId, + onCustomEvent: (event) => { + if (isCustomNodeEvent(event)) { + setCustomNodeEvents((prev) => [...prev, event]); + } + }, }); + const initializeEvents = customNodeEvents.filter( + (e) => e.nodeId === INITIALIZE_NODE_ID, + ); + const steps = mapCustomEventsToSteps(initializeEvents); + const allSuccess = + steps.length > 0 && steps.every((s) => s.status === "success"); + + let initStatus: "loading" | "generating" | "done" = "generating"; + if (allSuccess) { + initStatus = "done"; + } + + useEffect(() => { + const customInitEvents = getCustomNodeEventsFromMessages( + stream.messages, + INITIALIZE_NODE_ID, + ); + // If there are no custom init events found in messages, or we already have steps from custom events, return + if (!customInitEvents?.length || initializeEvents.length) { + return; + } + setCustomNodeEvents(customInitEvents); + }, [stream.messages]); + const streamJoined = useRef(false); useEffect(() => { if (!streamJoined.current && runId) { streamJoined.current = true; + // TODO: If the SDK changes go in, use this instead: + // stream.joinStream(runId, undefined, { streamMode: ["values", "messages", "custom"]}).catch(console.error); stream.joinStream(runId).catch(console.error); } }, [runId]); - const nonHumanMessages = stream.messages?.filter( - (m) => !isHumanMessageSDK(m), + // Filter out human & do not render messages + const filteredMessages = stream.messages?.filter( + (m) => + !isHumanMessageSDK(m) && + !(m.id && m.id.startsWith(DO_NOT_RENDER_ID_PREFIX)), ); // TODO: Need a better way to handle this. Not great like this... @@ -60,7 +128,15 @@ export function ActionsRenderer({ return (
- {nonHumanMessages?.map((m) => ( + {initializeEvents.length > 0 && steps.length > 0 && ( + + )} + {filteredMessages?.map((m) => ( >} diff --git a/apps/web/src/components/v2/terminal-input.tsx b/apps/web/src/components/v2/terminal-input.tsx index 9432cbbf..92d0b7d0 100644 --- a/apps/web/src/components/v2/terminal-input.tsx +++ b/apps/web/src/components/v2/terminal-input.tsx @@ -10,8 +10,6 @@ import { Button } from "../ui/button"; import { useStream } from "@langchain/langgraph-sdk/react"; import { useRouter } from "next/navigation"; import { useGitHubAppProvider } from "@/providers/GitHubApp"; -import { Message } from "@langchain/langgraph-sdk"; -import { useFileUpload } from "@/hooks/useFileUpload"; import { GraphState } from "@open-swe/shared/open-swe/types"; import { Base64ContentBlock, HumanMessage } from "@langchain/core/messages"; import { toast } from "sonner"; @@ -46,16 +44,17 @@ export function TerminalInput({ apiUrl, assistantId, reconnectOnMount: true, - threadId: null, - onThreadId: (id) => { - push(`/chat/${id}`); - setLoading(false); - setMessage(""); - setContentBlocks([]); - }, }); - const handleSend = () => { + const handleSend = async () => { + const assistantId = process.env.NEXT_PUBLIC_MANAGER_ASSISTANT_ID; + if (!assistantId) { + toast.error("No assistant ID found", { + richColors: true, + closeButton: true, + }); + return; + } if (!selectedRepository) { toast.error("Please select a repository first", { richColors: true, @@ -76,25 +75,34 @@ export function TerminalInput({ ], }); - stream.submit( - { - messages: [newHumanMessage], - targetRepository: selectedRepository, - }, - { - streamMode: ["values"], - optimisticValues: (prev) => ({ - ...prev, - messages: [...(prev.messages ?? []), newHumanMessage], - }), + try { + const newThreadId = uuidv4(); + const run = await stream.client.runs.create(newThreadId, assistantId, { + input: { + messages: [newHumanMessage], + targetRepository: selectedRepository, + }, config: { recursion_limit: 400, configurable: { ...getConfig(DEFAULT_CONFIG_KEY), }, }, - }, - ); + ifNotExists: "create", + streamResumable: true, + streamMode: ["values", "messages", "custom"], + }); + + // set session storage so the stream can be resumed after redirect. + sessionStorage.setItem(`lg:stream:${newThreadId}`, run.run_id); + push(`/chat/${newThreadId}`); + setMessage(""); + setContentBlocks([]); + } catch (e) { + console.error(e); + } finally { + setLoading(false); + } } }; @@ -108,7 +116,6 @@ export function TerminalInput({ return (
- {/* User@Host */} open-swe @ github diff --git a/apps/web/src/components/v2/thread-view.tsx b/apps/web/src/components/v2/thread-view.tsx index 1a001b35..fb3e31dc 100644 --- a/apps/web/src/components/v2/thread-view.tsx +++ b/apps/web/src/components/v2/thread-view.tsx @@ -16,6 +16,7 @@ import { PlannerGraphState } from "@open-swe/shared/open-swe/planner/types"; import { ActionsRenderer } from "./actions-renderer"; import { ThemeToggle } from "../theme-toggle"; import { HumanMessage } from "@langchain/core/messages"; +import { DO_NOT_RENDER_ID_PREFIX } from "@open-swe/shared/constants"; const PROGRAMMER_ASSISTANT_ID = process.env.NEXT_PUBLIC_PROGRAMMER_ASSISTANT_ID; const PLANNER_ASSISTANT_ID = process.env.NEXT_PUBLIC_PLANNER_ASSISTANT_ID; @@ -52,13 +53,26 @@ export function ThreadView({ id: uuidv4(), content: chatInput, }); - stream.submit({ - messages: [newHumanMessage], - }); + stream.submit( + { + messages: [newHumanMessage], + }, + { + streamResumable: true, + optimisticValues: (prev) => ({ + ...prev, + messages: [...(prev.messages ?? []), newHumanMessage], + }), + }, + ); setChatInput(""); } }; + const filteredMessages = stream.messages.filter((message) => { + return !message.id?.startsWith(DO_NOT_RENDER_ID_PREFIX); + }); + return (
{/* Header */} @@ -105,7 +119,7 @@ export function ThreadView({
{/* Chat Messages */}
- {stream.messages.map((message) => ( + {filteredMessages.map((message) => (
; -type StreamContextType = ReturnType; +type StreamContextType = ReturnType & { + customEvents: CustomNodeEvent[]; +}; const StreamContext = createContext(undefined); async function sleep(ms = 4000) { @@ -51,13 +57,17 @@ const StreamSession = ({ githubToken: string; }) => { const [threadId, setThreadId] = useQueryState("threadId"); - const { refreshThreads, setThreads } = useThreadsContext(); + const [customEvents, setCustomEvents] = useState([]); + const { refreshThreads } = useThreadsContext(); const streamValue = useTypedStream({ apiUrl, assistantId, reconnectOnMount: true, threadId: threadId ?? null, onCustomEvent: (event, options) => { + if (isCustomNodeEvent(event)) { + setCustomEvents((prev) => [...prev, event]); + } if (isUIMessage(event) || isRemoveUIMessage(event)) { options.mutate((prev) => { const ui = uiMessageReducer(prev.ui ?? [], event); @@ -75,7 +85,12 @@ const StreamSession = ({ }); return ( - + {children} ); diff --git a/packages/shared/src/constants.ts b/packages/shared/src/constants.ts index 06483e34..76c027e3 100644 --- a/packages/shared/src/constants.ts +++ b/packages/shared/src/constants.ts @@ -7,3 +7,5 @@ export const PLAN_INTERRUPT_ACTION_TITLE = "Approve/Edit Plan"; // Prefix the access token with `x-` so that it's included in requests to the LangGraph server. export const GITHUB_TOKEN_COOKIE = "x-github-access-token"; export const GITHUB_INSTALLATION_TOKEN_COOKIE = "x-github-installation-token"; + +export const DO_NOT_RENDER_ID_PREFIX = "do-not-render-"; diff --git a/packages/shared/src/open-swe/custom-node-events.ts b/packages/shared/src/open-swe/custom-node-events.ts new file mode 100644 index 00000000..f868edc6 --- /dev/null +++ b/packages/shared/src/open-swe/custom-node-events.ts @@ -0,0 +1,74 @@ +export type CustomNodeEvent = { + /** + * A UUID for the node the action is associated with. + */ + nodeId: string; + /** + * A UUID for the action the event is associated with. + */ + actionId: string; + action: string; + createdAt: string; + data: { + status: "pending" | "success" | "error" | "skipped"; + [key: string]: unknown; + }; +}; + +export function isCustomNodeEvent(event: unknown): event is CustomNodeEvent { + return ( + typeof event === "object" && + event !== null && + "nodeId" in event && + "actionId" in event && + "action" in event && + "data" in event && + "createdAt" in event + ); +} +export const INITIALIZE_NODE_ID = "initialize"; + +export const INIT_STEPS = [ + "Resuming Sandbox", + "Creating Sandbox", + "Cloning repository", + "Pulling latest changes", + "Configuring git user", + "Checking out branch", + "Generating codebase tree", +]; + +export type Step = { + name: string; + status: "waiting" | "generating" | "success" | "error" | "skipped"; + error?: string; +}; + +/** + * Maps custom events to step objects for UI rendering. Skipped steps are filtered out. + */ +export function mapCustomEventsToSteps(events: CustomNodeEvent[]) { + return INIT_STEPS.map((stepName) => { + const event = [...events] + .filter((e) => e.action === stepName) + .sort( + (a, b) => + new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime(), + )[0]; + if (!event) return { name: stepName, status: "waiting" as const }; + if (event.data.status === "skipped") + return { name: stepName, status: "skipped" as const }; + if (event.data.status === "pending") + return { name: stepName, status: "generating" as const }; + if (event.data.status === "success") + return { name: stepName, status: "success" as const }; + if (event.data.status === "error") + return { + name: stepName, + status: "error" as const, + error: + typeof event.data.error === "string" ? event.data.error : undefined, + }; + return { name: stepName, status: "waiting" as const }; + }).filter((step) => step.status !== "skipped"); +}