diff --git a/src/builders.spec.tsx b/src/builders.spec.tsx
index de7249d..a4922b7 100644
--- a/src/builders.spec.tsx
+++ b/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
hello world
;
+ },
+ }).provide({
+ actions: {
+ someAction: () => {},
+ },
+ });
+
+ const Root = buildRootComponent({ machine: xstateTreeMachine });
+ render();
+ });
});
describe("viewToMachine", () => {
diff --git a/src/builders.tsx b/src/builders.tsx
index 2219e02..4616a25 100644
--- a/src/builders.tsx
+++ b/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
+>(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;
}
/**
diff --git a/src/types.ts b/src/types.ts
index 320ddf4..0e7d931 100644
--- a/src/types.ts
+++ b/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
+ ) => XstateTreeMachine
+ : T[K];
+};
+
/**
* @public
*/
@@ -110,7 +127,12 @@ export type XstateTreeMachine<
TSelectorsOutput = ContextFrom,
TActionsOutput = Record,
TSlots extends readonly Slot[] = Slot[]
-> = TMachine &
+> = RepairProvideReturnType<
+ TMachine,
+ TSelectorsOutput,
+ TActionsOutput,
+ TSlots
+> &
XstateTreeMachineInjection<
TMachine,
TSelectorsOutput,