Implement "Run From Here" and fix some graph processor bugs

packages/app-executor/bin/executor.mts

@@ -47,7 +47,7 @@ const rivetDebugger = startDebuggerServer({
   port,
   allowGraphUpload: true,
   datasetProvider,
-  dynamicGraphRun: async ({ graphId, inputs, runToNodeIds, contextValues }) => {
+  dynamicGraphRun: async ({ graphId, inputs, runToNodeIds, contextValues, runFromNodeId }) => {
     console.log(`Running graph ${graphId} with inputs:`, inputs);
 
     const project = currentDebuggerState.uploadedProject;
@@ -146,6 +146,10 @@ const rivetDebugger = startDebuggerServer({
         processor.processor.runToNodeIds = runToNodeIds;
       }
 
+      if (runFromNodeId) {
+        processor.processor.runFromNodeId = runFromNodeId;
+      }
+
       await processor.run();
     } catch (err) {
       console.error(err);

packages/app/src/components/NodeCanvas.tsx

@@ -565,6 +565,8 @@ export const NodeCanvas: FC<NodeCanvasProps> = ({
     handleContextMenu(e);
   });
 
+  const lastRunPerNode = useRecoilValue(lastRunDataByNodeState);
+
   const hydratedContextMenuData = useMemo((): ContextMenuContext | null => {
     if (contextMenuData.data?.type.startsWith('node-')) {
       const nodeType = contextMenuData.data.type.replace('node-', '');
@@ -574,6 +576,7 @@ export const NodeCanvas: FC<NodeCanvasProps> = ({
         data: {
           nodeType,
           nodeId,
+          canRunFromHere: lastRunPerNode[nodeId] != null,
         },
       };
     }
@@ -582,7 +585,7 @@ export const NodeCanvas: FC<NodeCanvasProps> = ({
       type: 'blankArea',
       data: {},
     };
-  }, [contextMenuData]);
+  }, [contextMenuData, lastRunPerNode]);
 
   // Idk, before we were able to unmount the context menu, but safari be weird,
   // so we move it off screen instead
@@ -596,7 +599,6 @@ export const NodeCanvas: FC<NodeCanvasProps> = ({
   const pinnedNodes = useRecoilValue(pinnedNodesState);
 
   const nodeTypes = useNodeTypes();
-  const lastRunPerNode = useRecoilValue(lastRunDataByNodeState);
   const selectedProcessPagePerNode = useRecoilValue(selectedProcessPageNodesState);
 
   const isZoomedOut = canvasPosition.zoom < 0.4;

packages/app/src/hooks/useContextMenuConfiguration.ts

@@ -61,6 +61,7 @@ export function useContextMenuConfiguration() {
             contextType: type<{
               nodeType: string;
               nodeId: NodeId;
+              canRunFromHere: boolean;
             }>(),
             items: [
               {
@@ -98,6 +99,15 @@ export function useContextMenuConfiguration() {
                 label: 'Run to Here',
                 icon: PlayIcon,
               },
+              {
+                id: 'node-run-from-here',
+                label: 'Run from Here',
+                icon: PlayIcon,
+                conditional: (context) => {
+                  const { canRunFromHere } = context as { canRunFromHere: boolean };
+                  return canRunFromHere;
+                },
+              },
               {
                 id: 'node-delete',
                 label: 'Delete',

packages/app/src/hooks/useGraphBuilderContextMenuHandler.ts

@@ -111,6 +111,11 @@ export function useGraphBuilderContextMenuHandler() {
 
           tryRunGraph({ to: [nodeId] });
         })
+        .with('node-run-from-here', () => {
+          const { nodeId } = context.data as { nodeId: NodeId };
+
+          tryRunGraph({ from: nodeId });
+        })
         .with('node-copy', () => {
           const { nodeId } = context.data as { nodeId: NodeId };
           copyNodes(nodeId);

packages/app/src/hooks/useLocalExecutor.ts

@@ -8,6 +8,7 @@ import {
   type GraphOutputs,
   globalRivetNodeRegistry,
   type GraphId,
+  type Outputs,
 } from '@ironclad/rivet-core';
 import { produce } from 'immer';
 import { useRef } from 'react';
@@ -27,6 +28,7 @@ import { trivetState } from '../state/trivet';
 import { runTrivet } from '@ironclad/trivet';
 import { audioProvider, datasetProvider } from '../utils/globals';
 import { entries } from '../../../core/src/utils/typeSafety';
+import { type RunDataByNodeId, lastRunData, lastRunDataByNodeState } from '../state/dataFlow';
 
 export function useLocalExecutor() {
   const project = useRecoilValue(projectState);
@@ -43,6 +45,7 @@ export function useLocalExecutor() {
   const recordExecutions = useRecoilValue(recordExecutionsState);
   const projectData = useRecoilValue(projectDataState);
   const projectContext = useRecoilValue(projectContextState(project.metadata.id));
+  const lastRunData = useRecoilValue(lastRunDataByNodeState);
 
   function attachGraphEvents(processor: GraphProcessor) {
     processor.on('nodeStart', currentExecution.onNodeStart);
@@ -90,6 +93,7 @@ export function useLocalExecutor() {
       options: {
         graphId?: GraphId;
         to?: NodeId[];
+        from?: NodeId;
       } = {},
     ) => {
       try {
@@ -120,6 +124,11 @@ export function useLocalExecutor() {
           processor.runToNodeIds = options.to;
         }
 
+        if (options.from) {
+          preloadDependentDataForNode(processor, options.from, lastRunData);
+          processor.runFromNodeId = options.from;
+        }
+
         if (recordExecutions) {
           recorder.record(processor);
         }
@@ -257,3 +266,42 @@ export function useLocalExecutor() {
     tryRunTests,
   };
 }
+
+function preloadDependentDataForNode(processor: GraphProcessor, nodeId: NodeId, previousRunData: RunDataByNodeId) {
+  const dependencyNodes = processor.getDependencyNodesDeep(nodeId);
+
+  for (const dependencyNode of dependencyNodes) {
+    const dependencyNodeData = previousRunData[dependencyNode];
+
+    if (!dependencyNodeData) {
+      throw new Error(`Node ${dependencyNode} was not found in the previous run data, cannot continue preloading data`);
+    }
+
+    const firstExecution = dependencyNodeData[0];
+
+    if (!firstExecution?.data.outputData) {
+      throw new Error(
+        `Node ${dependencyNode} has no output data in the previous run data, cannot continue preloading data`,
+      );
+    }
+
+    const { outputData } = firstExecution.data;
+
+    // Convert back to DataValue from DataValueWithRefs
+    const outputDataWithoutRefs = Object.fromEntries(
+      Object.entries(outputData).map(([portId, dataValueWithRefs]) => {
+        if (dataValueWithRefs.type === 'image') {
+          throw new Error('Not implemented yed');
+        } else if (dataValueWithRefs.type === 'binary') {
+          throw new Error('Not implemented yed');
+        } else if (dataValueWithRefs.type === 'audio') {
+          throw new Error('Not implemented yed');
+        } else {
+          return [portId, dataValueWithRefs];
+        }
+      }),
+    ) as Outputs;
+
+    processor.preloadNodeData(dependencyNode, outputDataWithoutRefs);
+  }
+}

packages/app/src/hooks/useRemoteDebugger.ts

@@ -23,7 +23,6 @@ export function useRemoteDebugger(options: { onConnect?: () => void; onDisconnec
 
   const connectRef = useRef<((url: string) => void) | undefined>();
   const reconnectingTimeout = useRef<ReturnType<typeof setTimeout> | undefined>();
-  const selectedExecutor = useRecoilValue(selectedExecutorState);
 
   connectRef.current = (url: string) => {
     if (!url) {

packages/app/src/hooks/useRemoteExecutor.ts

@@ -7,6 +7,8 @@ import {
   serializeDatasets,
   type GraphId,
   type DataValue,
+  GraphProcessor,
+  type Outputs,
 } from '@ironclad/rivet-core';
 import { useCurrentExecution } from './useCurrentExecution';
 import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil';
@@ -25,6 +27,7 @@ import { pluginsState } from '../state/plugins';
 import { entries } from '../../../core/src/utils/typeSafety';
 import { selectedExecutorState } from '../state/execution';
 import { datasetProvider } from '../utils/globals';
+import { type RunDataByNodeId, lastRunDataByNodeState } from '../state/dataFlow';
 
 // TODO: This allows us to retrieve the GraphOutputs from the remote debugger.
 // If the remote debugger events had a unique ID for each run, this would feel a lot less hacky.
@@ -46,6 +49,7 @@ export function useRemoteExecutor() {
   const setUserInputQuestions = useSetRecoilState(userInputModalQuestionsState);
   const selectedExecutor = useRecoilValue(selectedExecutorState);
   const projectContext = useRecoilValue(projectContextState(project.metadata.id));
+  const lastRunData = useRecoilValue(lastRunDataByNodeState);
 
   const remoteDebugger = useRemoteDebugger({
     onDisconnect: () => {
@@ -119,7 +123,7 @@ export function useRemoteExecutor() {
     }
   });
 
-  const tryRunGraph = async (options: { to?: NodeId[]; graphId?: GraphId } = {}) => {
+  const tryRunGraph = async (options: { to?: NodeId[]; from?: NodeId; graphId?: GraphId } = {}) => {
     if (
       !remoteDebugger.remoteDebuggerState.started ||
       remoteDebugger.remoteDebuggerState.socket?.readyState !== WebSocket.OPEN
@@ -171,7 +175,21 @@ export function useRemoteExecutor() {
         {} as Record<string, DataValue>,
       );
 
-      remoteDebugger.send('run', { graphId: graphToRun, runToNodeIds: options.to, contextValues });
+      if (options.from) {
+        // Use a local graph processor to get dependency nodes instead of asking the remote debugger
+        const processor = new GraphProcessor(project, graph.metadata!.id!);
+        const dependencyNodes = processor.getDependencyNodesDeep(options.from);
+        const preloadData = getDependentDataForNodeForPreload(dependencyNodes, lastRunData);
+
+        remoteDebugger.send('preload', { nodeData: preloadData });
+      }
+
+      remoteDebugger.send('run', {
+        graphId: graphToRun,
+        runToNodeIds: options.to,
+        contextValues,
+        runFromNodeIds: options.from,
+      });
     } catch (e) {
       console.error(e);
     }
@@ -304,3 +322,44 @@ export function useRemoteExecutor() {
     tryRunTests,
   };
 }
+
+function getDependentDataForNodeForPreload(dependencyNodes: NodeId[], previousRunData: RunDataByNodeId) {
+  const preloadData: Record<NodeId, Outputs> = {};
+
+  for (const dependencyNode of dependencyNodes) {
+    const dependencyNodeData = previousRunData[dependencyNode];
+
+    if (!dependencyNodeData) {
+      throw new Error(`Node ${dependencyNode} was not found in the previous run data, cannot continue preloading data`);
+    }
+
+    const firstExecution = dependencyNodeData[0];
+
+    if (!firstExecution?.data.outputData) {
+      throw new Error(
+        `Node ${dependencyNode} has no output data in the previous run data, cannot continue preloading data`,
+      );
+    }
+
+    const { outputData } = firstExecution.data;
+
+    // Convert back to DataValue from DataValueWithRefs
+    const outputDataWithoutRefs = Object.fromEntries(
+      Object.entries(outputData).map(([portId, dataValueWithRefs]) => {
+        if (dataValueWithRefs.type === 'image') {
+          throw new Error('Not implemented yed');
+        } else if (dataValueWithRefs.type === 'binary') {
+          throw new Error('Not implemented yed');
+        } else if (dataValueWithRefs.type === 'audio') {
+          throw new Error('Not implemented yed');
+        } else {
+          return [portId, dataValueWithRefs];
+        }
+      }),
+    ) as Outputs;
+
+    preloadData[dependencyNode] = outputDataWithoutRefs;
+  }
+
+  return preloadData;
+}