diff --git a/packages/react/src/hooks/useOnSelectionChange.ts b/packages/react/src/hooks/useOnSelectionChange.ts
index 65cad6fd..17232895 100644
--- a/packages/react/src/hooks/useOnSelectionChange.ts
+++ b/packages/react/src/hooks/useOnSelectionChange.ts
@@ -1,10 +1,10 @@
 import { useEffect } from 'react';
 
 import { useStoreApi } from './useStore';
-import type { OnSelectionChangeFunc } from '../types';
+import type { OnSelectionChangeFunc, Node, Edge } from '../types';
 
-export type UseOnSelectionChangeOptions = {
-  onChange: OnSelectionChangeFunc;
+export type UseOnSelectionChangeOptions<NodeType extends Node = Node, EdgeType extends Edge = Edge> = {
+  onChange: OnSelectionChangeFunc<NodeType, EdgeType>;
 };
 
 /**
@@ -45,8 +45,10 @@ export type UseOnSelectionChangeOptions = {
  *
  * @remarks You need to memoize the passed `onChange` handler, otherwise the hook will not work correctly.
  */
-export function useOnSelectionChange({ onChange }: UseOnSelectionChangeOptions) {
-  const store = useStoreApi();
+export function useOnSelectionChange<NodeType extends Node = Node, EdgeType extends Edge = Edge>({
+  onChange,
+}: UseOnSelectionChangeOptions<NodeType, EdgeType>) {
+  const store = useStoreApi<NodeType, EdgeType>();
 
   useEffect(() => {
     const nextOnSelectionChangeHandlers = [...store.getState().onSelectionChangeHandlers, onChange];
diff --git a/packages/react/src/types/component-props.ts b/packages/react/src/types/component-props.ts
index b7c414f1..d558bc47 100644
--- a/packages/react/src/types/component-props.ts
+++ b/packages/react/src/types/component-props.ts
@@ -216,7 +216,7 @@ export interface ReactFlowProps<NodeType extends Node = Node, EdgeType extends E
   /** This event handler gets called when a user stops panning or zooming the viewport */
   onMoveEnd?: OnMoveEnd;
   /** This event handler gets called when a user changes group of selected elements in the flow */
-  onSelectionChange?: OnSelectionChangeFunc;
+  onSelectionChange?: OnSelectionChangeFunc<NodeType, EdgeType>;
   /** This event handler gets called when user scroll inside the pane */
   onPaneScroll?: (event?: WheelEvent) => void;
   /** This event handler gets called when user clicks inside the pane */
@@ -316,7 +316,7 @@ export interface ReactFlowProps<NodeType extends Node = Node, EdgeType extends E
    *
    * By setting this prop to null you can disable this functionality.
    * @default 'Meta' for macOS, "Ctrl" for other systems
-   * 
+   *
    */
   zoomActivationKeyCode?: KeyCode | null;
   /** Set this prop to make the flow snap to the grid */
diff --git a/packages/react/src/types/general.ts b/packages/react/src/types/general.ts
index 913ec634..8b78e77b 100644
--- a/packages/react/src/types/general.ts
+++ b/packages/react/src/types/general.ts
@@ -85,17 +85,19 @@ export type EdgeTypes = Record<
   >
 >;
 
-export type UnselectNodesAndEdgesParams = {
-  nodes?: Node[];
-  edges?: Edge[];
+export type UnselectNodesAndEdgesParams<NodeType extends Node = Node, EdgeType extends Edge = Edge> = {
+  nodes?: NodeType[];
+  edges?: EdgeType[];
 };
 
-export type OnSelectionChangeParams = {
-  nodes: Node[];
-  edges: Edge[];
+export type OnSelectionChangeParams<NodeType extends Node = Node, EdgeType extends Edge = Edge> = {
+  nodes: NodeType[];
+  edges: EdgeType[];
 };
 
-export type OnSelectionChangeFunc = (params: OnSelectionChangeParams) => void;
+export type OnSelectionChangeFunc<NodeType extends Node = Node, EdgeType extends Edge = Edge> = (
+  params: OnSelectionChangeParams<NodeType, EdgeType>
+) => void;
 
 export type FitViewParams<NodeType extends Node = Node> = FitViewParamsBase<NodeType>;
 
diff --git a/packages/react/src/types/store.ts b/packages/react/src/types/store.ts
index b060be36..e0ec76fa 100644
--- a/packages/react/src/types/store.ts
+++ b/packages/react/src/types/store.ts
@@ -134,7 +134,7 @@ export type ReactFlowStore<NodeType extends Node = Node, EdgeType extends Edge =
   onViewportChangeEnd?: OnViewportChange;
   onBeforeDelete?: OnBeforeDelete<NodeType, EdgeType>;
 
-  onSelectionChangeHandlers: OnSelectionChangeFunc[];
+  onSelectionChangeHandlers: OnSelectionChangeFunc<NodeType, EdgeType>[];
 
   ariaLiveMessage: string;
   autoPanOnConnect: boolean;
@@ -155,7 +155,7 @@ export type ReactFlowActions<NodeType extends Node, EdgeType extends Edge> = {
   updateNodeInternals: (updates: Map<string, InternalNodeUpdate>, params?: { triggerFitView: boolean }) => void;
   updateNodePositions: UpdateNodePositions;
   resetSelectedElements: () => void;
-  unselectNodesAndEdges: (params?: UnselectNodesAndEdgesParams) => void;
+  unselectNodesAndEdges: (params?: UnselectNodesAndEdgesParams<NodeType, EdgeType>) => void;
   addSelectedNodes: (nodeIds: string[]) => void;
   addSelectedEdges: (edgeIds: string[]) => void;
   setMinZoom: (minZoom: number) => void;