fix code node bug

Author: mariusandra

frontend/src/scenes/frame/panels/Diagram/CodeNode.tsx

@@ -1,5 +1,5 @@
 import { BindLogic, useActions, useValues } from 'kea'
-import { NodeProps, Handle, Position, NodeResizer } from 'reactflow'
+import { NodeProps, Handle, Position, NodeResizer, useUpdateNodeInternals } from 'reactflow'
 import { CodeNodeData } from '../../../../types'
 import clsx from 'clsx'
 import { diagramLogic } from './diagramLogic'
@@ -12,8 +12,9 @@ import { CodeArg } from './CodeArg'
 import { newNodePickerLogic } from './newNodePickerLogic'
 
 export function CodeNode({ id, isConnectable }: NodeProps<CodeNodeData>): JSX.Element {
+  const updateNodeInternals = useUpdateNodeInternals()
   const { frameId, sceneId } = useValues(diagramLogic)
-  const { updateNodeData, copyAppJSON, deleteApp } = useActions(diagramLogic)
+  const { updateNodeData, updateEdge, copyAppJSON, deleteApp } = useActions(diagramLogic)
   const appNodeLogicProps = { frameId, sceneId, nodeId: id }
   const { isSelected, node, nodeEdges } = useValues(appNodeLogic(appNodeLogicProps))
   const data: CodeNodeData = (node?.data as CodeNodeData) ?? ({ code: '' } satisfies CodeNodeData)
@@ -93,11 +94,17 @@ export function CodeNode({ id, isConnectable }: NodeProps<CodeNodeData>): JSX.El
                     <CodeArg
                       key={`${codeField.type}/${codeField.name}`}
                       codeArg={codeField}
-                      onChange={(value) =>
+                      onChange={(value) => {
                         updateNodeData(id, {
                           codeArgs: data.codeArgs?.map((c, j) => (i === j ? { ...c, ...value } : c)),
                         })
-                      }
+                        nodeEdges.forEach((edge) => {
+                          if (edge.target === id && edge.targetHandle === `codeField/${codeField.name}`) {
+                            updateEdge({ ...edge, targetHandle: `codeField/${value.name}` })
+                          }
+                        })
+                        updateNodeInternals(id)
+                      }}
                       onDelete={() => editCodeField(codeField.name, '')}
                     />
                   </div>

frontend/src/scenes/frame/panels/Diagram/diagramLogic.tsx

@@ -74,6 +74,7 @@ export const diagramLogic = kea<diagramLogicType>([
     fitDiagramView: true,
     keywordDropped: (keyword: string, type: string, position: XYPosition) => ({ keyword, type, position }),
     updateNodeData: (id: string, data: Record<string, any>) => ({ id, data }),
+    updateEdge: (edge: Edge) => ({ edge }),
     updateNodeConfig: (id: string, field: string, value: any) => ({ id, field, value }),
     copyAppJSON: (nodeId: string) => ({ nodeId }),
     deleteApp: (id: string) => ({ id }),
@@ -139,6 +140,10 @@ export const diagramLogic = kea<diagramLogicType>([
           const newEdges = state.filter((edge) => edge.source !== id && edge.target !== id)
           return equal(state, newEdges) ? state : newEdges
         },
+        updateEdge: (state, { edge }) => {
+          const newEdges = state.map((oldEdge) => (oldEdge.id === edge.id ? { ...oldEdge, ...edge } : oldEdge))
+          return equal(state, newEdges) ? state : newEdges
+        },
       },
     ],
     fitViewCounter: [0, { fitDiagramView: (state) => state + 1 }],

frontend/src/scenes/frame/panels/Diagram/newNodePickerLogic.tsx

@@ -505,9 +505,13 @@ export const newNodePickerLogic = kea<newNodePickerLogicType>([
         return
       }
 
+      // Dragged onto the canvas from a "+" codefield arg
       if (handleId === 'codeField/+') {
         const codeArgs = (values.nodes.find((node) => node.id === nodeId)?.data as CodeNodeData)?.codeArgs ?? []
-        let newArg = { name: keyword || getNewFieldName(codeArgs), type: type ?? 'string' } satisfies CodeArg
+        let newArg = {
+          name: keyword === '+' ? getNewFieldName(codeArgs) : keyword,
+          type: type ?? 'string',
+        } satisfies CodeArg
         actions.setNodes([
           ...values.nodes.map((node) =>
             node.id === nodeId ? { ...node, data: { ...node.data, codeArgs: [...codeArgs, newArg] } } : node

frontend/src/scenes/frame/panels/Scenes/ExpandedScene.tsx

@@ -7,6 +7,7 @@ import { TextArea } from '../../../../components/TextArea'
 import { TextInput } from '../../../../components/TextInput'
 import { Button } from '../../../../components/Button'
 import { controlLogic } from './controlLogic'
+import { panelsLogic } from '../panelsLogic'
 
 export interface ExpandedSceneProps {
   sceneId: string
@@ -17,6 +18,7 @@ export function ExpandedScene({ frameId, sceneId }: ExpandedSceneProps) {
   const { stateChanges, hasStateChanges, fields } = useValues(expandedSceneLogic({ frameId, sceneId }))
   const { states, sceneId: currentSceneId } = useValues(controlLogic({ frameId }))
   const { submitStateChanges, resetStateChanges } = useActions(expandedSceneLogic({ frameId, sceneId }))
+  const { editScene } = useActions(panelsLogic)
   const fieldCount = fields.length ?? 0
 
   const currentState = states[sceneId] ?? {}
@@ -26,9 +28,14 @@ export function ExpandedScene({ frameId, sceneId }: ExpandedSceneProps) {
       {fieldCount === 0 ? (
         <div className="space-y-2">
           <div>This scene does not export publicly controllable state.</div>
-          <Button onClick={submitStateChanges} color={sceneId !== currentSceneId ? 'primary' : 'secondary'}>
-            Activate scene
-          </Button>
+          <div className="flex items-center gap-2">
+            <Button onClick={submitStateChanges} color={sceneId !== currentSceneId ? 'primary' : 'secondary'}>
+              Activate scene
+            </Button>
+            <Button onClick={() => editScene(sceneId)} color="secondary">
+              Edit scene
+            </Button>
+          </div>
         </div>
       ) : (
         <Form logic={expandedSceneLogic} props={{ frameId, sceneId }} formKey="stateChanges" className="space-y-2">
@@ -100,6 +107,9 @@ export function ExpandedScene({ frameId, sceneId }: ExpandedSceneProps) {
               <Button onClick={() => resetStateChanges()} color="secondary">
                 Reset
               </Button>
+              <Button onClick={() => editScene(sceneId)} color="secondary">
+                Edit scene
+              </Button>
             </div>
           ) : null}
         </Form>