refactor(useConnection): return internal node, add node generic

packages/react/src/hooks/useConnection.ts

@@ -1,10 +1,10 @@
 import { shallow } from 'zustand/shallow';
-
-import { useStore } from './useStore';
-import type { ReactFlowStore } from '../types/store';
 import { ConnectionState, pointToRendererPoint } from '@xyflow/system';
 
-const selector = (s: ReactFlowStore): ConnectionState => {
+import { useStore } from './useStore';
+import type { InternalNode, Node, ReactFlowStore } from '../types';
+
+const selector = (s: ReactFlowStore) => {
   return s.connection.inProgress
     ? { ...s.connection, to: pointToRendererPoint(s.connection.to, s.transform) }
     : { ...s.connection };
@@ -15,6 +15,6 @@ const selector = (s: ReactFlowStore): ConnectionState => {
  * @public
  * @returns ConnectionState
  */
-export function useConnection(): ConnectionState {
-  return useStore(selector, shallow);
+export function useConnection<NodeType extends Node = Node>(): ConnectionState<InternalNode<NodeType>> {
+  return useStore(selector, shallow) as ConnectionState<InternalNode<NodeType>>;
 }

packages/react/src/types/edges.ts

@@ -16,7 +16,7 @@ import type {
   OnError,
 } from '@xyflow/system';
 
-import { EdgeTypes, Node } from '.';
+import { EdgeTypes, InternalNode, Node } from '.';
 
 export type EdgeLabelOptions = {
   label?: string | ReactNode;
@@ -190,10 +190,10 @@ export type SimpleBezierEdgeProps = EdgeComponentProps;
 
 export type OnReconnect<EdgeType extends Edge = Edge> = (oldEdge: EdgeType, newConnection: Connection) => void;
 
-export type ConnectionLineComponentProps = {
+export type ConnectionLineComponentProps<NodeType extends Node = Node> = {
   connectionLineStyle?: CSSProperties;
   connectionLineType: ConnectionLineType;
-  fromNode: Node;
+  fromNode: InternalNode<NodeType>;
   fromHandle: Handle;
   fromX: number;
   fromY: number;
@@ -202,8 +202,10 @@ export type ConnectionLineComponentProps = {
   fromPosition: Position;
   toPosition: Position;
   connectionStatus: 'valid' | 'invalid' | null;
-  toNode: Node | null;
+  toNode: InternalNode<NodeType> | null;
   toHandle: Handle | null;
 };
 
-export type ConnectionLineComponent = ComponentType<ConnectionLineComponentProps>;
+export type ConnectionLineComponent<NodeType extends Node = Node> = ComponentType<
+  ConnectionLineComponentProps<NodeType>
+>;

packages/system/src/types/general.ts

@@ -148,34 +148,30 @@ export const initialConnection: NoConnection = {
 export type NoConnection = {
   inProgress: false;
   isValid: null;
-
   from: null;
   fromHandle: null;
   fromPosition: null;
   fromNode: null;
-
   to: null;
   toHandle: null;
   toPosition: null;
   toNode: null;
 };
-
-export type ConnectionInProgress = {
+export type ConnectionInProgress<NodeType extends InternalNodeBase = InternalNodeBase> = {
   inProgress: true;
   isValid: boolean | null;
-
   from: XYPosition;
   fromHandle: Handle;
   fromPosition: Position;
-  fromNode: NodeBase;
-
+  fromNode: NodeType;
   to: XYPosition;
   toHandle: Handle | null;
   toPosition: Position;
-  toNode: NodeBase | null;
+  toNode: NodeType | null;
 };
-
-export type ConnectionState = ConnectionInProgress | NoConnection;
+export type ConnectionState<NodeType extends InternalNodeBase = InternalNodeBase> =
+  | ConnectionInProgress<NodeType>
+  | NoConnection;
 
 export type UpdateConnection = (params: ConnectionState) => void;
 

packages/system/src/xyhandle/XYHandle.ts

@@ -104,7 +104,7 @@ function onPointerDown(
     from,
     fromHandle,
     fromPosition: fromHandle.position,
-    fromNode: fromNodeInternal.internals.userNode,
+    fromNode: fromNodeInternal,
 
     to: position,
     toHandle: null,
@@ -163,7 +163,7 @@ function onPointerDown(
           : position,
       toHandle: result.toHandle,
       toPosition: isValid && result.toHandle ? result.toHandle.position : oppositePosition[fromHandle.position],
-      toNode: result.toHandle ? nodeLookup.get(result.toHandle.nodeId)!.internals.userNode : null,
+      toNode: result.toHandle ? nodeLookup.get(result.toHandle.nodeId)! : null,
     };
 
     // we don't want to trigger an update when the connection