Fix ast issue

package.json

@@ -2,7 +2,7 @@
 	"name": "react-router-devtools",
 	"description": "Devtools for React Router - debug, trace, find hydration errors, catch bugs and inspect server/client data with react-router-devtools",
 	"author": "Alem Tuzlak",
-	"version": "5.0.5",
+	"version": "5.0.6",
 	"license": "MIT",
 	"keywords": [
 		"react-router",

src/vite/plugin.tsx

@@ -251,8 +251,7 @@ export const reactRouterDevTools: (args?: ReactRouterViteConfig) => Plugin[] = (
 					process.rdt_port = server.config.server.port ?? 5173
 					port = process.rdt_port
 				})
-				//@ts-ignore - vite 5/6 compat
-				const channel = server.hot.channels.find((channel) => channel.name === "ws") ?? server.environments?.client.hot
+				const channel = server.hot
 				const editor = args?.editor ?? DEFAULT_EDITOR_CONFIG
 				const openInEditor = async (path: string | undefined, lineNum: string | undefined) => {
 					if (!path) {
@@ -268,9 +267,7 @@ export const reactRouterDevTools: (args?: ReactRouterViteConfig) => Plugin[] = (
 						}
 						if (routine === "request-event") {
 							unusedEvents.set(parsedData.id + parsedData.startTime, parsedData)
-							for (const client of server.hot.channels) {
-								client.send("request-event", JSON.stringify(parsedData))
-							}
+							server.hot.send("request-event", JSON.stringify(parsedData))
 
 							return
 						}
@@ -291,9 +288,7 @@ export const reactRouterDevTools: (args?: ReactRouterViteConfig) => Plugin[] = (
 								routeInfo.set(id, { loader: [], action: [data] })
 							}
 						}
-						for (const client of server.hot.channels) {
-							client.send("route-info", JSON.stringify({ type, data }))
-						}
+						server.hot.send("route-info", JSON.stringify({ type, data }))
 					})
 				)
 

src/vite/utils/data-functions-augment.test.ts

@@ -92,8 +92,30 @@ describe("transform", () => {
 			)
 			const expected = removeWhitespace(`
 			import { withLoaderWrapper as _withLoaderWrapper   } from "react-router-devtools/server";
-      import { loader  } from "./loader.js";
-			export { loader as _loader };
+      import { loader as _loader  } from "./loader.js";
+			export const loader = _withLoaderWrapper(_loader, "test");
+		`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
+		it("should wrap the loader export when it's imported from another file and exported and used", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+			import {  loader } from "./loader.js";
+			const test = () => {
+			  const data = loader();
+			}
+			export { loader };
+			`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+			import { withLoaderWrapper as _withLoaderWrapper   } from "react-router-devtools/server";
+      import { loader as _loader  } from "./loader.js";
+			const test = () => {
+			  const data = _loader();
+			};
 			export const loader = _withLoaderWrapper(_loader, "test");
 		`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
@@ -217,8 +239,30 @@ describe("transform", () => {
 			)
 			const expected = removeWhitespace(`
 				import { withClientLoaderWrapper as _withClientLoaderWrapper   } from "react-router-devtools/client";
-				import { clientLoader  } from "./client-loader.js";
-				export { clientLoader as _clientLoader };
+				import { clientLoader as _clientLoader } from "./client-loader.js";
+				export const clientLoader = _withClientLoaderWrapper(_clientLoader, "test");
+			`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
+		it("should wrap the client loader export when it's re-exported from another file and used by other code", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+				import { clientLoader } from "./client-loader.js";
+				const test = () => {
+					const data = clientLoader();
+				}
+				export { clientLoader };
+				`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+				import { withClientLoaderWrapper as _withClientLoaderWrapper   } from "react-router-devtools/client";
+				import { clientLoader as _clientLoader } from "./client-loader.js";
+				const test = () => {
+					const data = _clientLoader();
+				};
 				export const clientLoader = _withClientLoaderWrapper(_clientLoader, "test");
 			`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
@@ -235,8 +279,30 @@ describe("transform", () => {
 			)
 			const expected = removeWhitespace(`
 				import { withClientLoaderWrapper as _withClientLoaderWrapper   } from "react-router-devtools/client";
-				import { clientLoader  } from "./client-loader.js";
-				export { clientLoader as _clientLoader };
+				import { clientLoader as _clientLoader } from "./client-loader.js";
+				export const clientLoader = _withClientLoaderWrapper(_clientLoader, "test");
+			`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
+		it("should wrap the client loader export when it's imported from another file and exported and used by other code", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+				import { clientLoader } from "./client-loader.js";
+				const test = () => {
+					const data = clientLoader();
+				}
+				export { clientLoader };
+				`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+				import { withClientLoaderWrapper as _withClientLoaderWrapper   } from "react-router-devtools/client";
+				import { clientLoader as _clientLoader } from "./client-loader.js";
+				const test = () => {
+					const data = _clientLoader();
+				};
 				export const clientLoader = _withClientLoaderWrapper(_clientLoader, "test");
 			`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
@@ -373,8 +439,30 @@ describe("transform", () => {
 			)
 			const expected = removeWhitespace(`
 			import { withActionWrapper as _withActionWrapper   } from "react-router-devtools/server";
-      import { action  } from "./action.js";
-			export { action as _action };
+      import { action as _action } from "./action.js";
+			export const action = _withActionWrapper(_action, "test");
+		`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
+		it("should wrap the action export when it's imported from another file and exported and used by other code", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+			import {  action } from "./action.js";
+			const test = () => {
+				const data = action();
+			}
+			export { action };
+			`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+			import { withActionWrapper as _withActionWrapper   } from "react-router-devtools/server";
+      import { action as _action } from "./action.js";
+			const test = () => {
+				const data = _action();
+			};
 			export const action = _withActionWrapper(_action, "test");
 		`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
@@ -478,8 +566,31 @@ describe("transform", () => {
 			)
 			const expected = removeWhitespace(`
 			import { withClientActionWrapper as _withClientActionWrapper   } from "react-router-devtools/client";
-			import { clientAction  } from "./client-action.js";
-			export { clientAction as _clientAction };
+			import { clientAction as _clientAction  } from "./client-action.js";
+			export const clientAction = _withClientActionWrapper(_clientAction, "test");
+		`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
+		it("should transform the client action export when it's re-exported from another file and keep it working if used somewhere", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+			import { clientAction } from "./client-action.js";
+
+			const test = () => {
+				const data = clientAction();
+			}
+			export { clientAction };
+			`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+			import { withClientActionWrapper as _withClientActionWrapper   } from "react-router-devtools/client";
+			import { clientAction as _clientAction  } from "./client-action.js";
+			const test = () => {
+				const data = _clientAction();
+			};
 			export const clientAction = _withClientActionWrapper(_clientAction, "test");
 		`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
@@ -512,13 +623,77 @@ describe("transform", () => {
 			)
 			const expected = removeWhitespace(`
 			import { withClientActionWrapper as _withClientActionWrapper   } from "react-router-devtools/client";
-      import { clientAction  } from "./client-action.js";
-			export { clientAction as _clientAction };
+      import { clientAction as _clientAction } from "./client-action.js";
 			export const clientAction = _withClientActionWrapper(_clientAction, "test");
 		`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
 		})
 
+		it("should transform the client action export when it's imported from another file and exported and used by other code", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+			import { clientAction } from "./client-action.js";
+			const test = () => {
+				const data = clientAction();
+			}
+			export { clientAction };
+			`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+			import { withClientActionWrapper as _withClientActionWrapper   } from "react-router-devtools/client";
+      import { clientAction as _clientAction } from "./client-action.js";
+			const test = () => {
+				const data = _clientAction();
+			};
+			export const clientAction = _withClientActionWrapper(_clientAction, "test");
+		`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
+		it("should transform the client action export when it's imported from another file and exported and keep the export around if more things are exported", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+			import { clientAction } from "./client-action.js";
+			const test = () => {
+				const data = clientAction();
+			}
+			export { clientAction, test };
+			`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+			import { withClientActionWrapper as _withClientActionWrapper   } from "react-router-devtools/client";
+      import { clientAction as _clientAction } from "./client-action.js";
+			const test = () => {
+				const data = _clientAction();
+			};
+			export { test };
+			export const clientAction = _withClientActionWrapper(_clientAction, "test");
+		`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+		it("should transform the client action export when it's imported from another file and exported and keep the export around if more things are exported", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+			import { withClientActionContextWrapper as _withClientActionContextWrapper   } from "react-router-devtools/context";
+      import { clientAction as _clientAction } from "./client-action.js";
+			export const clientAction = _withClientActionContextWrapper(_clientAction, "test");
+			`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+			import { withClientActionWrapper as _withClientActionWrapper   } from "react-router-devtools/client";
+			import { withClientActionContextWrapper as _withClientActionContextWrapper   } from "react-router-devtools/context";
+      import { clientAction as _clientAction } from "./client-action.js";
+			export const clientAction = _withClientActionWrapper(_withClientActionContextWrapper(_clientAction, "test"), "test");
+		`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
 		it("should wrap the clientAction export when it's exported via export { clientAction } and declared within the file", () => {
 			const result = augmentDataFetchingFunctions(
 				`

src/vite/utils/data-functions-augment.ts

@@ -10,6 +10,8 @@ const ALL_EXPORTS = [...SERVER_COMPONENT_EXPORTS, ...CLIENT_COMPONENT_EXPORTS]
 const transform = (ast: ParseResult<Babel.File>, routeId: string) => {
 	const serverHocs: Array<[string, Babel.Identifier]> = []
 	const clientHocs: Array<[string, Babel.Identifier]> = []
+	const imports: Array<[string, Babel.Identifier]> = []
+
 	function getServerHocId(path: NodePath, hocName: string) {
 		const uid = path.scope.generateUidIdentifier(hocName)
 		const hasHoc = serverHocs.find(([name]) => name === hocName)
@@ -35,6 +37,26 @@ const transform = (ast: ParseResult<Babel.File>, routeId: string) => {
 
 	const importDeclarations: Babel.ImportDeclaration[] = []
 	trav(ast, {
+		ImportDeclaration(path) {
+			const specifiers = path.node.specifiers
+			for (const specifier of specifiers) {
+				if (!t.isImportSpecifier(specifier) || !t.isIdentifier(specifier.imported)) {
+					continue
+				}
+				const name = specifier.imported.name
+				if (!ALL_EXPORTS.includes(name)) {
+					continue
+				}
+				const isReimported = specifier.local.name !== name
+				const uniqueName = isReimported ? specifier.local : path.scope.generateUidIdentifier(name)
+				imports.push([name, uniqueName])
+				specifier.local = uniqueName
+				// Replace the import specifier with a new one
+				if (!isReimported) {
+					path.scope.rename(name, uniqueName.name)
+				}
+			}
+		},
 		ExportDeclaration(path) {
 			if (path.isExportNamedDeclaration()) {
 				const decl = path.get("declaration")
@@ -148,14 +170,13 @@ const transform = (ast: ParseResult<Babel.File>, routeId: string) => {
 					}
 				} else {
 					transformations.push(() => {
-						const uniqueName = path.scope.generateUidIdentifier(name).name
-						path.replaceWith(
-							t.exportNamedDeclaration(
-								null,
-								[t.exportSpecifier(t.identifier(name), t.identifier(uniqueName))],
-								path.node.source
-							)
+						const existingImport = imports.find(([existingName]) => existingName === name)
+						const uniqueName = existingImport?.[1].name ?? path.scope.generateUidIdentifier(name).name
+
+						const remainingSpecifiers = path.node.specifiers.filter(
+							(exportSpecifier) => !(t.isIdentifier(exportSpecifier.exported) && exportSpecifier.exported.name === name)
 						)
+						path.replaceWith(t.exportNamedDeclaration(null, remainingSpecifiers, path.node.source))
 
 						// Insert the wrapped export after the modified export statement
 						path.insertAfter(
@@ -169,6 +190,10 @@ const transform = (ast: ParseResult<Babel.File>, routeId: string) => {
 								[]
 							)
 						)
+						const newRemainingSpecifiers = path.node.specifiers.length
+						if (newRemainingSpecifiers === 0) {
+							path.remove()
+						}
 					})
 				}
 			}