fix(builders): patch provide to function correctly

Due to how provide works it loses the _xstateTree property that is attached to xstate-tree machines, making them unusable.

It also fixes the return type to return an XstateTreeMachine instead of a StateMachine

src/builders.spec.tsx

@@ -53,6 +53,32 @@ describe("xstate-tree builders", () => {
         }}
       />;
     });
+
+    it("repairs the provide function to not lose the _xstateTree property and return an XstateTreeMachine", () => {
+      const machine = setup({
+        actions: {
+          someAction: () => {},
+        },
+      }).createMachine({
+        initial: "idle",
+        states: {
+          idle: {},
+        },
+      });
+
+      const xstateTreeMachine = createXStateTreeMachine(machine, {
+        View() {
+          return <div>hello world</div>;
+        },
+      }).provide({
+        actions: {
+          someAction: () => {},
+        },
+      });
+
+      const Root = buildRootComponent({ machine: xstateTreeMachine });
+      render(<Root />);
+    });
   });
 
   describe("viewToMachine", () => {

src/builders.tsx

@@ -54,7 +54,23 @@ export function createXStateTreeMachine<
     slots: (options.slots ?? []) as any,
   };
 
-  return machineWithMeta;
+  return fixProvideLosingXstateTreeMeta(machineWithMeta);
+}
+
+function fixProvideLosingXstateTreeMeta<
+  T extends XstateTreeMachine<any, any, any, any>
+>(machine: T): T {
+  const originalProvide = machine.provide.bind(machine);
+  (machine as any).provide = (impl: any) => {
+    const result = originalProvide(impl) as T;
+
+    result._xstateTree = machine._xstateTree;
+    fixProvideLosingXstateTreeMeta(result);
+
+    return result;
+  };
+
+  return machine;
 }
 
 /**

src/lazy.tsx

@@ -1,11 +1,11 @@
 import React from "react";
-import { AnyStateMachine, assign, setup, fromPromise, InputFrom } from "xstate";
+import { assign, setup, fromPromise, InputFrom } from "xstate";
 
 import { createXStateTreeMachine } from "./builders";
 import { singleSlot } from "./slots";
 import { AnyXstateTreeMachine } from "./types";
 
-type Options<TStateMachine extends AnyStateMachine> = {
+type Options<TStateMachine extends AnyXstateTreeMachine> = {
   /**
    * Displayed while the promise is resolving, defaults to returning null
    */

src/types.ts

@@ -102,6 +102,23 @@ export type XstateTreeMachineInjection<
   >;
 };
 
+/**
+ * Repairs the return type of the `provide` function on XstateTreeMachines to correctly return
+ * an XstateTreeMachine type instead of an xstate StateMachine
+ */
+type RepairProvideReturnType<
+  T extends AnyStateMachine,
+  TSelectorsOutput,
+  TActionsOutput,
+  TSlots extends readonly Slot[]
+> = {
+  [K in keyof T]: K extends "provide"
+    ? (
+        ...args: Parameters<T[K]>
+      ) => XstateTreeMachine<T, TSelectorsOutput, TActionsOutput, TSlots>
+    : T[K];
+};
+
 /**
  * @public
  */
@@ -110,7 +127,12 @@ export type XstateTreeMachine<
   TSelectorsOutput = ContextFrom<TMachine>,
   TActionsOutput = Record<never, string>,
   TSlots extends readonly Slot[] = Slot[]
-> = TMachine &
+> = RepairProvideReturnType<
+  TMachine,
+  TSelectorsOutput,
+  TActionsOutput,
+  TSlots
+> &
   XstateTreeMachineInjection<
     TMachine,
     TSelectorsOutput,

xstate-tree.api.md

@@ -380,10 +380,11 @@ export type XstateTreeHistory<T = unknown> = History_2<{
     previousUrl?: string;
 }>;
 
+// Warning: (ae-forgotten-export) The symbol "RepairProvideReturnType" needs to be exported by the entry point index.d.ts
 // Warning: (ae-incompatible-release-tags) The symbol "XstateTreeMachine" is marked as @public, but its signature references "XstateTreeMachineInjection" which is marked as @internal
 //
 // @public (undocumented)
-export type XstateTreeMachine<TMachine extends AnyStateMachine, TSelectorsOutput = ContextFrom<TMachine>, TActionsOutput = Record<never, string>, TSlots extends readonly Slot[] = Slot[]> = TMachine & XstateTreeMachineInjection<TMachine, TSelectorsOutput, TActionsOutput, TSlots>;
+export type XstateTreeMachine<TMachine extends AnyStateMachine, TSelectorsOutput = ContextFrom<TMachine>, TActionsOutput = Record<never, string>, TSlots extends readonly Slot[] = Slot[]> = RepairProvideReturnType<TMachine, TSelectorsOutput, TActionsOutput, TSlots> & XstateTreeMachineInjection<TMachine, TSelectorsOutput, TActionsOutput, TSlots>;
 
 // Warning: (ae-internal-missing-underscore) The name "XstateTreeMachineInjection" should be prefixed with an underscore because the declaration is marked as @internal
 //
@@ -400,8 +401,8 @@ export type XstateTreeMachineStateSchemaV2<TMachine extends AnyStateMachine, TSe
 // src/routing/createRoute/createRoute.ts:285:19 - (ae-forgotten-export) The symbol "MergeRouteTypes" needs to be exported by the entry point index.d.ts
 // src/routing/createRoute/createRoute.ts:285:19 - (ae-forgotten-export) The symbol "ResolveZodType" needs to be exported by the entry point index.d.ts
 // src/routing/createRoute/createRoute.ts:322:9 - (ae-forgotten-export) The symbol "RouteRedirect" needs to be exported by the entry point index.d.ts
-// src/types.ts:142:3 - (ae-incompatible-release-tags) The symbol "canHandleEvent" is marked as @public, but its signature references "CanHandleEvent" which is marked as @internal
-// src/types.ts:143:3 - (ae-incompatible-release-tags) The symbol "inState" is marked as @public, but its signature references "MatchesFrom" which is marked as @internal
+// src/types.ts:164:3 - (ae-incompatible-release-tags) The symbol "canHandleEvent" is marked as @public, but its signature references "CanHandleEvent" which is marked as @internal
+// src/types.ts:165:3 - (ae-incompatible-release-tags) The symbol "inState" is marked as @public, but its signature references "MatchesFrom" which is marked as @internal
 
 // (No @packageDocumentation comment for this package)