From 8512d344ef88423a4a9292713b2c0a5b17e60e86 Mon Sep 17 00:00:00 2001
From: peterkogo <peter@xyflow.com>
Date: Tue, 21 Oct 2025 14:26:41 +0200
Subject: [PATCH] use new createContext utility for better type safety

---
 .../lib/components/EdgeLabel/EdgeLabel.svelte | 17 +++++----
 .../EdgeReconnectAnchor.svelte                | 10 +++---
 .../components/EdgeWrapper/EdgeWrapper.svelte | 10 +++---
 .../src/lib/components/Handle/Handle.svelte   |  9 ++---
 .../components/NodeWrapper/NodeWrapper.svelte | 17 +++++----
 .../src/lib/hooks/derivedWarning.svelte.ts    | 24 -------------
 .../lib/hooks/useNodeConnections.svelte.ts    |  4 +--
 packages/svelte/src/lib/hooks/useStore.ts     |  5 ---
 .../hooks/useUpdateNodeInternals.svelte.ts    |  4 +--
 .../plugins/NodeResizer/ResizeControl.svelte  | 15 +++++---
 .../plugins/NodeToolbar/NodeToolbar.svelte    |  9 +++--
 packages/svelte/src/lib/store/context.ts      | 36 +++++++++++++++++++
 12 files changed, 89 insertions(+), 71 deletions(-)
 delete mode 100644 packages/svelte/src/lib/hooks/derivedWarning.svelte.ts
 create mode 100644 packages/svelte/src/lib/store/context.ts

diff --git a/packages/svelte/src/lib/components/EdgeLabel/EdgeLabel.svelte b/packages/svelte/src/lib/components/EdgeLabel/EdgeLabel.svelte
index b01ce307..ab48ef0d 100644
--- a/packages/svelte/src/lib/components/EdgeLabel/EdgeLabel.svelte
+++ b/packages/svelte/src/lib/components/EdgeLabel/EdgeLabel.svelte
@@ -1,11 +1,11 @@
 <script lang="ts">
-  import { getContext } from 'svelte';
-
+  import { getEdgeIdContext } from '$lib/store/context';
   import { hideOnSSR, portal } from '$lib/actions/portal';
   import { useStore } from '$lib/store';
-  import type { EdgeLabelProps } from './types';
   import { toPxString } from '$lib/utils';
 
+  import type { EdgeLabelProps } from './types';
+
   let {
     x = 0,
     y = 0,
@@ -19,13 +19,12 @@
   }: EdgeLabelProps = $props();
 
   const store = useStore();
-  const edgeId = getContext<string>('svelteflow__edge_id');
 
-  if (!edgeId) {
-    throw new Error('EdgeLabel must be used within an edge');
-  }
+  const edgeId = getEdgeIdContext('EdgeLabel must be used within a Custom Edge component');
 
-  let z = $derived(store.visible.edges.get(edgeId)?.zIndex);
+  let z = $derived.by(() => {
+    return store.visible.edges.get(edgeId)?.zIndex;
+  });
 </script>
 
 <div
@@ -40,7 +39,7 @@
   style:z-index={z}
   tabindex="-1"
   onclick={() => {
-    if (selectEdgeOnClick) store.handleEdgeSelection(edgeId);
+    if (selectEdgeOnClick && edgeId) store.handleEdgeSelection(edgeId);
   }}
   {...rest}
 >
diff --git a/packages/svelte/src/lib/components/EdgeReconnectAnchor/EdgeReconnectAnchor.svelte b/packages/svelte/src/lib/components/EdgeReconnectAnchor/EdgeReconnectAnchor.svelte
index 81f0ab2d..2a4975df 100644
--- a/packages/svelte/src/lib/components/EdgeReconnectAnchor/EdgeReconnectAnchor.svelte
+++ b/packages/svelte/src/lib/components/EdgeReconnectAnchor/EdgeReconnectAnchor.svelte
@@ -1,8 +1,8 @@
 <script lang="ts">
   import { useStore } from '$lib/store';
+  import { getEdgeIdContext } from '$lib/store/context';
   import type { Edge } from '$lib/types';
   import { XYHandle, type HandleType, type OnConnectStart } from '@xyflow/system';
-  import { getContext } from 'svelte';
   import { EdgeLabel } from '../EdgeLabel';
   import type { EdgeReconnectAnchorProps } from './types';
 
@@ -19,11 +19,9 @@
 
   const store = useStore();
 
-  let edgeId: string | undefined = getContext('svelteflow__edge_id');
-
-  if (!edgeId) {
-    throw new Error('EdgeReconnectAnchor must be used within an Edge component');
-  }
+  const edgeId = getEdgeIdContext(
+    'EdgeReconnectAnchor must be used within a Custom Edge component'
+  );
 
   const onPointerDown = (event: PointerEvent) => {
     if (event.button !== 0) {
diff --git a/packages/svelte/src/lib/components/EdgeWrapper/EdgeWrapper.svelte b/packages/svelte/src/lib/components/EdgeWrapper/EdgeWrapper.svelte
index 08dd4eb2..d86ec944 100644
--- a/packages/svelte/src/lib/components/EdgeWrapper/EdgeWrapper.svelte
+++ b/packages/svelte/src/lib/components/EdgeWrapper/EdgeWrapper.svelte
@@ -1,8 +1,8 @@
 <script lang="ts" generics="NodeType extends Node = Node, EdgeType extends Edge = Edge">
-  import { setContext } from 'svelte';
-
   import { elementSelectionKeys, getMarkerId } from '@xyflow/system';
 
+  import { setEdgeIdContext } from '$lib/store/context';
+
   import { BezierEdgeInternal } from '$lib/components/edges';
 
   import type { Node, EdgeLayouted, Edge, EdgeEvents } from '$lib/types';
@@ -22,6 +22,7 @@
   } & EdgeEvents<EdgeType> = $props();
 
   let {
+    id,
     source,
     target,
     sourceX,
@@ -51,12 +52,11 @@
     ariaLabel
   } = $derived(edge);
 
+  setEdgeIdContext(edge.id);
+
   // svelte-ignore non_reactive_update
   let edgeRef: SVGGElement | null = null;
 
-  const { id } = edge;
-  setContext('svelteflow__edge_id', id);
-
   let selectable = $derived(_selectable ?? store.elementsSelectable);
   let focusable = $derived(_focusable ?? store.edgesFocusable);
 
diff --git a/packages/svelte/src/lib/components/Handle/Handle.svelte b/packages/svelte/src/lib/components/Handle/Handle.svelte
index eb52320e..15f811b1 100644
--- a/packages/svelte/src/lib/components/Handle/Handle.svelte
+++ b/packages/svelte/src/lib/components/Handle/Handle.svelte
@@ -1,5 +1,4 @@
 <script lang="ts">
-  import { getContext } from 'svelte';
   import {
     Position,
     XYHandle,
@@ -16,8 +15,8 @@
 
   import { useStore } from '$lib/store';
 
-  import type { ConnectableContext } from '../NodeWrapper/types';
   import type { HandleProps } from './types';
+  import { getNodeConnectableContext, getNodeIdContext } from '$lib/store/context';
 
   let {
     id: handleId = null,
@@ -35,8 +34,10 @@
     ...rest
   }: HandleProps = $props();
 
-  const nodeId = getContext<string>('svelteflow__node_id');
-  const isConnectableContext = getContext<ConnectableContext>('svelteflow__node_connectable');
+  const nodeId = getNodeIdContext('Handle must be used within a Custom Node component');
+  const isConnectableContext = getNodeConnectableContext(
+    'Handle must be used within a Custom Node component'
+  );
 
   let isTarget = $derived(type === 'target');
   let isConnectable = $derived(
diff --git a/packages/svelte/src/lib/components/NodeWrapper/NodeWrapper.svelte b/packages/svelte/src/lib/components/NodeWrapper/NodeWrapper.svelte
index 7aec74fa..0dbb5875 100644
--- a/packages/svelte/src/lib/components/NodeWrapper/NodeWrapper.svelte
+++ b/packages/svelte/src/lib/components/NodeWrapper/NodeWrapper.svelte
@@ -1,5 +1,5 @@
 <script lang="ts" generics="NodeType extends Node = Node, EdgeType extends Edge = Edge">
-  import { setContext, onDestroy } from 'svelte';
+  import { onDestroy } from 'svelte';
   import {
     elementSelectionKeys,
     errorMessages,
@@ -10,13 +10,15 @@
   } from '@xyflow/system';
 
   import drag from '$lib/actions/drag';
-  import DefaultNode from '$lib/components/nodes/DefaultNode.svelte';
-
-  import type { ConnectableContext, NodeWrapperProps } from './types';
-  import type { Node, Edge, NodeEvents } from '$lib/types';
+  import { setNodeConnectableContext, setNodeIdContext } from '$lib/store/context';
   import { arrowKeyDiffs, toPxString } from '$lib/utils';
   import { ARIA_NODE_DESC_KEY } from '../A11yDescriptions';
 
+  import type { Node, Edge, NodeEvents } from '$lib/types';
+  import type { ConnectableContext, NodeWrapperProps } from './types';
+
+  import DefaultNode from '$lib/components/nodes/DefaultNode.svelte';
+
   let {
     store = $bindable(),
     node,
@@ -95,8 +97,9 @@
       return connectable;
     }
   };
-  setContext('svelteflow__node_connectable', connectableContext);
-  setContext('svelteflow__node_id', id);
+
+  setNodeIdContext(id);
+  setNodeConnectableContext(connectableContext);
 
   if (process.env.NODE_ENV === 'development') {
     $effect(() => {
diff --git a/packages/svelte/src/lib/hooks/derivedWarning.svelte.ts b/packages/svelte/src/lib/hooks/derivedWarning.svelte.ts
deleted file mode 100644
index c676f4b1..00000000
--- a/packages/svelte/src/lib/hooks/derivedWarning.svelte.ts
+++ /dev/null
@@ -1,24 +0,0 @@
-import { key } from '$lib/store';
-import type { StoreContext } from '$lib/store/types';
-import { getContext } from 'svelte';
-
-/**
- * Warns the user that they should use $derived() when calling a hook.
- * This is not neccessarry when the hook is called inside a child of <SvelteFlowFlow />,
- * however exceptions can be made if you don't want to return a closure.
- * @param functionName - The name of the function that is being called
- * @param force - If true, the warning will be shown regardless if child of <SvelteFlowFlow />
- */
-export function derivedWarning(functionName: string) {
-  const storeContext = getContext<StoreContext>(key);
-
-  if (!storeContext) {
-    throw new Error(
-      `In order to use ${functionName}() you need to wrap your component in a <SvelteFlowProvider />`
-    );
-  }
-
-  if (storeContext.provider && typeof window === 'object' && !$effect.tracking()) {
-    throw new Error(`Use $derived(${functionName}()) to receive updates when values change.`);
-  }
-}
diff --git a/packages/svelte/src/lib/hooks/useNodeConnections.svelte.ts b/packages/svelte/src/lib/hooks/useNodeConnections.svelte.ts
index b73091dc..c27f6532 100644
--- a/packages/svelte/src/lib/hooks/useNodeConnections.svelte.ts
+++ b/packages/svelte/src/lib/hooks/useNodeConnections.svelte.ts
@@ -8,7 +8,7 @@ import {
 } from '@xyflow/system';
 
 import { useStore } from '$lib/store';
-import { getContext } from 'svelte';
+import { getNodeIdContext } from '$lib/store/context';
 
 type UseNodeConnectionsParams = {
   id?: string;
@@ -42,7 +42,7 @@ export function useNodeConnections({
 }: UseNodeConnectionsParams = {}) {
   const { edges, connectionLookup } = $derived(useStore());
 
-  const contextNodeId = getContext<string>('svelteflow__node_id');
+  const contextNodeId = getNodeIdContext();
   const nodeId = id ?? contextNodeId;
 
   let connectionMaps: { previous: ConnectionMap; next: ConnectionMap } = {
diff --git a/packages/svelte/src/lib/hooks/useStore.ts b/packages/svelte/src/lib/hooks/useStore.ts
index 11c89e7e..732093e6 100644
--- a/packages/svelte/src/lib/hooks/useStore.ts
+++ b/packages/svelte/src/lib/hooks/useStore.ts
@@ -2,7 +2,6 @@ import { getContext } from 'svelte';
 import type { StoreContext, SvelteFlowStore } from '../store/types';
 
 import { key } from '../store';
-import { derivedWarning } from './derivedWarning.svelte';
 import type { Node, Edge } from '$lib/types';
 
 export function useStore<
@@ -17,9 +16,5 @@ export function useStore<
     );
   }
 
-  if (process.env.NODE_ENV === 'development') {
-    derivedWarning('useStore');
-  }
-
   return storeContext.getStore();
 }
diff --git a/packages/svelte/src/lib/hooks/useUpdateNodeInternals.svelte.ts b/packages/svelte/src/lib/hooks/useUpdateNodeInternals.svelte.ts
index 9dfe560b..e5febd36 100644
--- a/packages/svelte/src/lib/hooks/useUpdateNodeInternals.svelte.ts
+++ b/packages/svelte/src/lib/hooks/useUpdateNodeInternals.svelte.ts
@@ -1,6 +1,6 @@
 /* eslint-disable svelte/prefer-svelte-reactivity */
 import { useStore } from '$lib/store';
-import { getContext } from 'svelte';
+import { getNodeIdContext } from '$lib/store/context';
 
 /**
  * When you programmatically add or remove handles to a node or update a node's
@@ -14,7 +14,7 @@ import { getContext } from 'svelte';
  */
 export function useUpdateNodeInternals(): (nodeId?: string | string[]) => void {
   const { domNode, updateNodeInternals } = $derived(useStore());
-  const nodeId = getContext('svelteflow__node_id') as string | undefined;
+  const nodeId = getNodeIdContext();
 
   // @todo: do we want to add this to system?
   const updateInternals = (id?: string | string[]) => {
diff --git a/packages/svelte/src/lib/plugins/NodeResizer/ResizeControl.svelte b/packages/svelte/src/lib/plugins/NodeResizer/ResizeControl.svelte
index 8c6874e9..f53fa7a7 100644
--- a/packages/svelte/src/lib/plugins/NodeResizer/ResizeControl.svelte
+++ b/packages/svelte/src/lib/plugins/NodeResizer/ResizeControl.svelte
@@ -1,6 +1,8 @@
 <script lang="ts">
-  import { getContext, onMount } from 'svelte';
+  import { onMount } from 'svelte';
+
   import { useStore } from '$lib/store';
+  import { getNodeIdContext } from '$lib/store/context';
   import {
     XYResizer,
     ResizeControlVariant,
@@ -9,6 +11,7 @@
     type XYResizerChange,
     type XYResizerChildChange
   } from '@xyflow/system';
+
   import type { ResizeControlProps } from './types';
   import type { Node } from '$lib/types';
 
@@ -34,10 +37,14 @@
   }: ResizeControlProps = $props();
 
   const store = useStore();
+  const contextNodeId = getNodeIdContext();
 
-  let id = $derived(
-    typeof nodeId === 'string' ? nodeId : getContext<string>('svelteflow__node_id')
-  );
+  let id = $derived(typeof nodeId === 'string' ? nodeId : contextNodeId);
+
+  // svelte-ignore state_referenced_locally
+  if (!id) {
+    throw new Error('Either pass a nodeId or use within a Custom Node component');
+  }
 
   let resizeControlRef: HTMLDivElement;
   let resizer: XYResizerInstance | null = $state(null);
diff --git a/packages/svelte/src/lib/plugins/NodeToolbar/NodeToolbar.svelte b/packages/svelte/src/lib/plugins/NodeToolbar/NodeToolbar.svelte
index 49cfe3a1..55889b38 100644
--- a/packages/svelte/src/lib/plugins/NodeToolbar/NodeToolbar.svelte
+++ b/packages/svelte/src/lib/plugins/NodeToolbar/NodeToolbar.svelte
@@ -1,9 +1,9 @@
 <script lang="ts">
-  import { getContext } from 'svelte';
   import { Position, getNodeToolbarTransform } from '@xyflow/system';
 
-  import { hideOnSSR, portal } from '$lib/actions/portal';
   import { useStore } from '$lib/store';
+  import { getNodeIdContext } from '$lib/store/context';
+  import { hideOnSSR, portal } from '$lib/actions/portal';
   import { useSvelteFlow } from '$lib/hooks/useSvelteFlow.svelte';
 
   import type { InternalNode } from '$lib/types';
@@ -22,7 +22,7 @@
   const store = useStore();
 
   const { getNodesBounds } = useSvelteFlow();
-  const contextNodeId = getContext<string>('svelteflow__node_id');
+  const contextNodeId = getNodeIdContext();
 
   let toolbarNodes: InternalNode[] = $derived.by(() => {
     // eslint-disable-next-line @typescript-eslint/no-unused-expressions
@@ -30,6 +30,9 @@
     const nodeIds = Array.isArray(nodeId) ? nodeId : [nodeId ?? contextNodeId];
 
     return nodeIds.reduce<InternalNode[]>((res, nodeId) => {
+      if (!nodeId) {
+        throw new Error('Either pass a nodeId or use within a Custom Node component');
+      }
       const node = store.nodeLookup.get(nodeId);
 
       if (node) {
diff --git a/packages/svelte/src/lib/store/context.ts b/packages/svelte/src/lib/store/context.ts
new file mode 100644
index 00000000..79bd41d6
--- /dev/null
+++ b/packages/svelte/src/lib/store/context.ts
@@ -0,0 +1,36 @@
+import type { ConnectableContext } from '$lib/components/NodeWrapper/types';
+import { setContext, getContext, hasContext } from 'svelte';
+import type { StoreContext } from './types';
+
+/**
+ * Creates a type-safe context getter and setter pair.
+ * Extended from Svelte's official createContext pattern.
+ * - When called with an error message string, it throws if the context is not set
+ * - When called without arguments, it returns the context value or undefined
+ */
+function createContext<T>(): [
+  {
+    (errorMessage: string): T;
+    (): T | undefined;
+  },
+  (context: T) => T
+] {
+  const key = {};
+
+  return [
+    (errorMessage?: string) => {
+      if (errorMessage && !hasContext(key)) {
+        throw new Error(errorMessage);
+      }
+
+      return getContext(key);
+    },
+    (context) => setContext(key, context)
+  ];
+}
+
+export const [getNodeIdContext, setNodeIdContext] = createContext<string>();
+export const [getNodeConnectableContext, setNodeConnectableContext] =
+  createContext<ConnectableContext>();
+
+export const [getEdgeIdContext, setEdgeIdContext] = createContext<string>();