assistant-ui 0.0.122 → 0.0.124

package/src/codemods/v0-12/primitive-if-to-aui-if.ts

@@ -1,4 +1,5 @@
 import { createTransformer } from "../utils/createTransformer";
+import { resolveBinding } from "../utils/resolveBinding";
 
 type ConditionFragment = {
   expression: string;
@@ -185,37 +186,92 @@ const buildConditionString = (
 
 const migratePrimitiveIfToAuiIf = createTransformer(
   ({ j, root, markAsChanged }) => {
-    let needsAuiIfImport = false;
+    const importedPrimitives = new Map<
+      string,
+      { name: string; declaration: any; binding: any }
+    >();
+    const auiIfImports: { source: string; binding: any }[] = [];
+    const usedNames = new Set<string>();
+    root
+      .find(j.Identifier)
+      .forEach((path: any) => usedNames.add(path.value.name));
+    root
+      .find(j.JSXIdentifier)
+      .forEach((path: any) => usedNames.add(path.value.name));
 
-    // Track which primitive namespaces are imported
-    const importedPrimitives = new Set<string>();
     root.find(j.ImportDeclaration).forEach((path: any) => {
-      const source = path.value.source.value;
-      if (typeof source === "string" && source.startsWith("@assistant-ui/")) {
-        path.value.specifiers?.forEach((specifier: any) => {
-          if (j.ImportSpecifier.check(specifier)) {
-            const name = String(
-              specifier.local?.name ?? specifier.imported.name,
-            );
-            if (primitiveMap[name] || fixedConditionMap[name]) {
-              importedPrimitives.add(name);
-            }
-          }
-        });
-      }
+      const declaration = path.value;
+      const source = declaration.source.value;
+      if (
+        typeof source !== "string" ||
+        !source.startsWith("@assistant-ui/") ||
+        declaration.importKind === "type"
+      )
+        return;
+
+      declaration.specifiers?.forEach((specifier: any) => {
+        if (
+          !j.ImportSpecifier.check(specifier) ||
+          !j.Identifier.check(specifier.imported) ||
+          !j.Identifier.check(specifier.local) ||
+          (specifier as { importKind?: string }).importKind === "type"
+        )
+          return;
+        const name = specifier.imported.name;
+        const binding = specifier.local;
+        if (primitiveMap[name] || fixedConditionMap[name]) {
+          importedPrimitives.set(binding.name, { name, declaration, binding });
+        }
+        if (name === "AuiIf") auiIfImports.push({ source, binding });
+      });
     });
 
     if (importedPrimitives.size === 0) return;
 
+    const getPrimitive = (path: any, localName: string) => {
+      const primitive = importedPrimitives.get(localName);
+      return primitive &&
+        resolveBinding(j, path, localName) === primitive.binding
+        ? primitive
+        : undefined;
+    };
+
+    const getAuiIfName = (path: any, declaration: any) => {
+      const source = declaration.source.value;
+      const existing = auiIfImports.find(
+        (entry) =>
+          entry.source === source &&
+          !/^[a-z]/.test(entry.binding.name) &&
+          resolveBinding(j, path, entry.binding.name) === entry.binding,
+      );
+      if (existing) return existing.binding.name;
+
+      let name = "AuiIf";
+      for (let suffix = 1; usedNames.has(name); suffix++)
+        name = `AuiIf${suffix}`;
+      usedNames.add(name);
+      const binding = j.identifier(name);
+      declaration.specifiers.push(
+        j.importSpecifier(j.identifier("AuiIf"), binding),
+      );
+      auiIfImports.push({ source, binding });
+      return name;
+    };
+
     // Opening and closing tags are rewritten together per element, so a
     // skipped element can never be left with a mismatched closing tag.
-    const convertElementToAuiIf = (elementPath: any, conditionBody: string) => {
+    const convertElementToAuiIf = (
+      elementPath: any,
+      conditionBody: string,
+      declaration: any,
+    ) => {
+      const localName = getAuiIfName(elementPath, declaration);
       const arrowFnAst = j(`(s) => ${conditionBody}`)
         .find(j.ArrowFunctionExpression)
         .paths()[0]!.value;
 
       const opening = elementPath.value.openingElement;
-      opening.name = j.jsxIdentifier("AuiIf");
+      opening.name = j.jsxIdentifier(localName);
       opening.attributes = [
         j.jsxAttribute(
           j.jsxIdentifier("condition"),
@@ -223,10 +279,9 @@ const migratePrimitiveIfToAuiIf = createTransformer(
         ),
       ];
       if (elementPath.value.closingElement) {
-        elementPath.value.closingElement.name = j.jsxIdentifier("AuiIf");
+        elementPath.value.closingElement.name = j.jsxIdentifier(localName);
       }
 
-      needsAuiIfImport = true;
       markAsChanged();
     };
 
@@ -237,17 +292,17 @@ const migratePrimitiveIfToAuiIf = createTransformer(
       if (!j.JSXIdentifier.check(name.object)) return;
       if (!j.JSXIdentifier.check(name.property)) return;
 
-      const primitiveName = name.object.name as string;
+      const primitive = getPrimitive(path, name.object.name);
+      if (!primitive) return;
       const propertyName = name.property.name as string;
-      const conditionBody = fixedConditionMap[primitiveName]?.[propertyName];
+      const conditionBody = fixedConditionMap[primitive.name]?.[propertyName];
       if (!conditionBody) return;
-      if (!importedPrimitives.has(primitiveName)) return;
 
       // Only transform if there are no props (other than children, which are implicit)
       const attrs: any[] = path.value.openingElement.attributes || [];
       if (attrs.length > 0) return;
 
-      convertElementToAuiIf(path, conditionBody);
+      convertElementToAuiIf(path, conditionBody, primitive.declaration);
     });
 
     // Process JSX elements: <ThreadPrimitive.If ...> → <AuiIf condition={...}>
@@ -260,10 +315,10 @@ const migratePrimitiveIfToAuiIf = createTransformer(
       if (!j.JSXIdentifier.check(name.property)) return;
       if (name.property.name !== "If") return;
 
-      const primitiveName = name.object.name;
-      const propMap = primitiveMap[primitiveName];
+      const primitive = getPrimitive(path, name.object.name);
+      if (!primitive) return;
+      const propMap = primitiveMap[primitive.name];
       if (!propMap) return;
-      if (!importedPrimitives.has(primitiveName)) return;
 
       // Extract props
       const attrs: any[] = path.value.openingElement.attributes || [];
@@ -304,37 +359,12 @@ const migratePrimitiveIfToAuiIf = createTransformer(
       // If we couldn't map all props, skip this element
       if (hasUnknownProp || fragments.length === 0) return;
 
-      convertElementToAuiIf(path, buildConditionString(j, fragments));
+      convertElementToAuiIf(
+        path,
+        buildConditionString(j, fragments),
+        primitive.declaration,
+      );
     });
-
-    // Add AuiIf import if needed
-    if (needsAuiIfImport) {
-      let hasAuiIfImport = false;
-      let assistantUiImport: any = null;
-
-      root.find(j.ImportDeclaration).forEach((path: any) => {
-        const source = path.value.source.value;
-        if (typeof source === "string" && source.startsWith("@assistant-ui/")) {
-          assistantUiImport = path;
-          path.value.specifiers?.forEach((specifier: any) => {
-            if (
-              j.ImportSpecifier.check(specifier) &&
-              (specifier.imported.name === "AuiIf" ||
-                specifier.local?.name === "AuiIf")
-            ) {
-              hasAuiIfImport = true;
-            }
-          });
-        }
-      });
-
-      if (!hasAuiIfImport && assistantUiImport) {
-        assistantUiImport.value.specifiers.push(
-          j.importSpecifier(j.identifier("AuiIf")),
-        );
-        markAsChanged();
-      }
-    }
   },
 );
 

package/src/codemods/v0-8/__tests__/ui-package-split.test.ts

@@ -0,0 +1,83 @@
+import { describe, expect, it } from "vitest";
+import jscodeshift, { type API } from "jscodeshift";
+import transform from "../ui-package-split";
+
+const j = jscodeshift.withParser("tsx");
+const api: API = { jscodeshift: j, j, stats: () => {}, report: () => {} };
+const apply = (source: string) =>
+  transform({ path: "app.tsx", source }, api, {});
+
+describe("ui package split", () => {
+  it("preserves an aliased moved binding beside the same exported name", () => {
+    const output =
+      apply(`import { Thread as MigratedThread } from "@assistant-ui/react";
+import { Thread as ExistingThread } from "@assistant-ui/react-ui";
+const thread = MigratedThread;`)!;
+    const imports = j(output)
+      .find(j.ImportDeclaration, {
+        source: { value: "@assistant-ui/react-ui" },
+      })
+      .nodes();
+
+    expect(
+      imports.flatMap((declaration) =>
+        declaration.specifiers?.map((specifier) => specifier.local?.name),
+      ),
+    ).toEqual(expect.arrayContaining(["MigratedThread", "ExistingThread"]));
+    expect(output).toContain("const thread = MigratedThread;");
+  });
+
+  it("preserves a moved binding beside a namespace import", () => {
+    const output = apply(`import { Thread } from "@assistant-ui/react";
+import * as ui from "@assistant-ui/react-ui";
+const thread = Thread;`)!;
+    const imports = j(output)
+      .find(j.ImportDeclaration, {
+        source: { value: "@assistant-ui/react-ui" },
+      })
+      .nodes();
+
+    expect(
+      imports.flatMap((declaration) =>
+        declaration.specifiers?.map((specifier) => specifier.local?.name),
+      ),
+    ).toEqual(expect.arrayContaining(["Thread", "ui"]));
+    expect(output).toContain("const thread = Thread;");
+  });
+
+  it("keeps a moved value import out of an existing type declaration", () => {
+    const output = apply(`import { Composer } from "@assistant-ui/react";
+import type { ThreadConfig } from "@assistant-ui/react-ui";
+const composer = Composer;`)!;
+    const imports = j(output)
+      .find(j.ImportDeclaration, {
+        source: { value: "@assistant-ui/react-ui" },
+      })
+      .nodes();
+    const valueImport = imports.find((declaration) =>
+      declaration.specifiers?.some(
+        (specifier) => specifier.local?.name === "Composer",
+      ),
+    );
+
+    expect(valueImport).toBeDefined();
+    expect(valueImport?.importKind).not.toBe("type");
+    expect(output).toContain("const composer = Composer;");
+  });
+
+  it("keeps the declaration-level type kind when splitting an import", () => {
+    const output = apply(
+      'import type { ThreadConfig, AssistantRuntime } from "@assistant-ui/react";',
+    )!;
+    const declarations = j(output).find(j.ImportDeclaration).nodes();
+
+    expect(declarations).toHaveLength(2);
+    expect(declarations.map((declaration) => declaration.importKind)).toEqual([
+      "type",
+      "type",
+    ]);
+    expect(declarations.map((declaration) => declaration.source.value)).toEqual(
+      ["@assistant-ui/react", "@assistant-ui/react-ui"],
+    );
+  });
+});

package/src/codemods/v0-8/ui-package-split.ts

@@ -33,72 +33,34 @@ const migrateAssistantUI = createTransformer(({ j, root, markAsChanged }) => {
     "@assistant-ui/react",
     "@assistant-ui/react-markdown",
   ];
-  const movedSpecifiers: any[] = [];
-  let lastMigratedImportPath: any = null;
 
   root
     .find(j.ImportDeclaration)
-    .filter((path: any) => sourcesToMigrate.includes(path.value.source.value))
-    .forEach((path: any) => {
-      let hadMigratedSpecifiers = false;
-      const remainingSpecifiers: any[] = [];
-      path.value.specifiers.forEach((specifier: any) => {
-        if (
+    .filter((path) =>
+      sourcesToMigrate.includes(path.node.source.value as string),
+    )
+    .forEach((path) => {
+      const specifiers = path.node.specifiers ?? [];
+      const moved = specifiers.filter(
+        (specifier) =>
           j.ImportSpecifier.check(specifier) &&
-          reactUIExports.includes(specifier.imported.name as string)
-        ) {
-          movedSpecifiers.push(specifier);
-          hadMigratedSpecifiers = true;
-        } else {
-          remainingSpecifiers.push(specifier);
-        }
-      });
-      if (hadMigratedSpecifiers) {
-        lastMigratedImportPath = path;
-      }
-      if (remainingSpecifiers.length === 0) {
-        j(path).remove();
-        markAsChanged();
-      } else if (remainingSpecifiers.length !== path.value.specifiers.length) {
-        path.value.specifiers = remainingSpecifiers;
-        markAsChanged();
-      }
-    });
-
-  if (movedSpecifiers.length > 0) {
-    const existingReactUIImport = root.find(j.ImportDeclaration, {
-      source: { value: "@assistant-ui/react-ui" },
-    });
-    if (existingReactUIImport.size() > 0) {
-      existingReactUIImport.forEach((path: any) => {
-        movedSpecifiers.forEach((specifier: any) => {
-          if (
-            !path.value.specifiers.some(
-              (s: any) => s.imported.name === specifier.imported.name,
-            )
-          ) {
-            path.value.specifiers.push(specifier);
-          }
-        });
-      });
-    } else {
-      const newImport = j.importDeclaration(
-        movedSpecifiers,
-        j.literal("@assistant-ui/react-ui"),
+          reactUIExports.includes(specifier.imported.name as string),
       );
-      if (lastMigratedImportPath) {
-        j(lastMigratedImportPath).insertAfter(newImport);
+      if (moved.length === 0) return;
+
+      const destination = j.literal("@assistant-ui/react-ui");
+      if (moved.length === specifiers.length) {
+        path.node.source = destination;
       } else {
-        const firstImport = root.find(j.ImportDeclaration).at(0);
-        if (firstImport.size() > 0) {
-          firstImport.insertBefore(newImport);
-        } else {
-          root.get().node.program.body.unshift(newImport);
-        }
+        path.node.specifiers = specifiers.filter(
+          (specifier) => !moved.includes(specifier),
+        );
+        const declaration = j.importDeclaration(moved, destination);
+        declaration.importKind = path.node.importKind ?? "value";
+        j(path).insertAfter(declaration);
       }
-    }
-    markAsChanged();
-  }
+      markAsChanged();
+    });
 
   const cssReplacements: Record<string, string> = {
     "@assistant-ui/react/styles/index.css":

package/src/codemods/v0-9/edge-package-split.test.ts

@@ -0,0 +1,100 @@
+import { describe, expect, it } from "vitest";
+import jscodeshift, { type API } from "jscodeshift";
+import transform from "./edge-package-split";
+
+const j = jscodeshift.withParser("tsx");
+const api: API = { jscodeshift: j, j, stats: () => {}, report: () => {} };
+const apply = (source: string) =>
+  transform({ path: "app.tsx", source }, api, {});
+
+describe.each(["@assistant-ui/react", "@assistant-ui/react-edge"])(
+  "edge migration from %s",
+  (source) => {
+    it.each([
+      'import { useChatRuntime as existingRuntime } from "@assistant-ui/react-ai-sdk";',
+      'import * as sdk from "@assistant-ui/react-ai-sdk";',
+      'import { useChatRuntime as first } from "@assistant-ui/react-ai-sdk";\nimport { useChatRuntime as second } from "@assistant-ui/react-ai-sdk";',
+      'import type { CoreMessage } from "@assistant-ui/react-ai-sdk";',
+    ])("preserves the moved local binding beside %s", (existing) => {
+      const output = apply(`import { useChatRuntime } from "${source}";
+${existing}
+const runtime = useChatRuntime;`)!;
+
+      expect(output).toContain(existing);
+      const imports = j(output)
+        .find(j.ImportDeclaration, {
+          source: { value: "@assistant-ui/react-ai-sdk" },
+        })
+        .nodes();
+      const binding = imports.filter((declaration) =>
+        declaration.specifiers?.some(
+          (specifier) => specifier.local?.name === "useChatRuntime",
+        ),
+      );
+      expect(binding).toHaveLength(1);
+      expect(binding[0]!.importKind).not.toBe("type");
+      expect(output).toContain("const runtime = useChatRuntime;");
+      expect(apply(output)).toBeNull();
+    });
+
+    it("keeps declaration-level type imports type-only", () => {
+      const output = apply(`import type { CoreMessage } from "${source}";
+type Message = CoreMessage;`)!;
+      const declaration = j(output).find(j.ImportDeclaration).nodes()[0]!;
+      expect(declaration.source.value).toBe("@assistant-ui/react-ai-sdk");
+      expect(declaration.importKind).toBe("type");
+      expect(output).toContain("type Message = CoreMessage;");
+    });
+  },
+);
+
+it("keeps unrelated bindings and specifier-level type imports", () => {
+  const output =
+    apply(`import { AssistantRuntimeProvider, type CoreMessage, useChatRuntime as useRuntime } from "@assistant-ui/react";
+export { AssistantRuntimeProvider, useRuntime };
+export type { CoreMessage };`)!;
+  const root = j(output);
+  const declarations = root.find(j.ImportDeclaration).nodes();
+  const retained = declarations.find(
+    (declaration) => declaration.source.value === "@assistant-ui/react",
+  )!;
+  expect(
+    retained.specifiers?.map((specifier) => specifier.local?.name),
+  ).toEqual(["AssistantRuntimeProvider"]);
+  const moved = root.find(j.ImportSpecifier).nodes();
+  expect(
+    moved.find((specifier) => specifier.local?.name === "CoreMessage"),
+  ).toMatchObject({ importKind: "type" });
+  expect(
+    moved.find((specifier) => specifier.local?.name === "useRuntime")?.imported
+      .name,
+  ).toBe("useChatRuntime");
+  expect(apply(output)).toBeNull();
+});
+
+it("keeps both halves of a split type declaration type-only", () => {
+  const output = apply(
+    'import type { AssistantRuntime, CoreMessage } from "@assistant-ui/react";',
+  )!;
+  const declarations = j(output).find(j.ImportDeclaration).nodes();
+  expect(declarations).toHaveLength(2);
+  expect(declarations.map((declaration) => declaration.importKind)).toEqual([
+    "type",
+    "type",
+  ]);
+  expect(declarations.map((declaration) => declaration.source.value)).toEqual([
+    "@assistant-ui/react",
+    "@assistant-ui/react-ai-sdk",
+  ]);
+});
+
+it.each(["edge/converters", "dangerous-in-browser/converters"])(
+  "continues migrating the %s subpath",
+  (subpath) => {
+    expect(
+      apply(
+        `export { toCoreMessages } from "@assistant-ui/react/${subpath}";\nimport { fromCoreMessages } from "@assistant-ui/react/${subpath}";`,
+      ),
+    ).toContain(`from "@assistant-ui/react-ai-sdk/${subpath}"`);
+  },
+);