Synchronize tabs and persist to localstorage

packages/ui/app/src/atoms/codeGroupTab.ts

@@ -0,0 +1,18 @@
+import { useAtomValue } from "jotai";
+import { atomWithStorage } from "jotai/utils";
+
+export const FERN_CODE_GROUP_TAB = atomWithStorage<{ groupId: string; value: string }[]>("code-group-tab-id", []);
+FERN_CODE_GROUP_TAB.debugLabel = "FERN_CODE_GROUP_TAB";
+
+interface Group {
+    groupId: string;
+    value: string;
+}
+
+export function useCodeGroup({ groupId }: { groupId?: string }): { selectedGroup?: Group; groupIds: Group[] } {
+    const groupIds = useAtomValue(FERN_CODE_GROUP_TAB);
+    if (groupId) {
+        return { selectedGroup: groupIds.find((pair) => pair.groupId === groupId), groupIds };
+    }
+    return { selectedGroup: undefined, groupIds };
+}

packages/ui/app/src/atoms/index.ts

@@ -1,5 +1,6 @@
 export * from "./apis";
 export * from "./auth";
+export * from "./codeGroupTab";
 export * from "./docs";
 export * from "./files";
 export * from "./flags";

packages/ui/app/src/mdx/components/code/CodeGroup.tsx

@@ -1,8 +1,9 @@
 import { CopyToClipboardButton } from "@fern-ui/components";
 import * as Tabs from "@radix-ui/react-tabs";
 import clsx from "clsx";
-import { useState } from "react";
-import { useFeatureFlags } from "../../../atoms";
+import { useAtom } from "jotai";
+import { useEffect, useState } from "react";
+import { FERN_CODE_GROUP_TAB, useCodeGroup, useFeatureFlags } from "../../../atoms";
 import { HorizontalOverflowMask } from "../../../commons/HorizontalOverflowMask";
 import { FernSyntaxHighlighter, FernSyntaxHighlighterProps } from "../../../syntax-highlighting/FernSyntaxHighlighter";
 
@@ -13,12 +14,15 @@ export declare namespace CodeGroup {
 
     export interface Props {
         items: Item[];
+        groupId?: string;
     }
 }
 
-export const CodeGroup: React.FC<React.PropsWithChildren<CodeGroup.Props>> = ({ items }) => {
+export const CodeGroup: React.FC<React.PropsWithChildren<CodeGroup.Props>> = ({ items, groupId }) => {
     const { isDarkCodeEnabled } = useFeatureFlags();
-    const [selectedTabIndex, setSelectedTabIndex] = useState(0);
+    const [selectedTabIndex, setSelectedTabIndex] = useState("0");
+    const [_, setSelectedTab] = useAtom(FERN_CODE_GROUP_TAB);
+    const { selectedGroup, groupIds } = useCodeGroup({ groupId });
 
     const containerClass = clsx(
         "after:ring-card-border bg-card relative mt-4 first:mt-0 mb-6 flex w-full min-w-0 max-w-full flex-col rounded-lg shadow-sm after:pointer-events-none after:absolute after:inset-0 after:rounded-[inherit] after:ring-1 after:ring-inset after:content-['']",
@@ -27,6 +31,12 @@ export const CodeGroup: React.FC<React.PropsWithChildren<CodeGroup.Props>> = ({
         },
     );
 
+    useEffect(() => {
+        if (!!selectedGroup) {
+            setSelectedTabIndex(selectedGroup.value);
+        }
+    }, [selectedGroup]);
+
     if (items.length === 1 && items[0] != null) {
         return (
             <div className={containerClass}>
@@ -50,8 +60,18 @@ export const CodeGroup: React.FC<React.PropsWithChildren<CodeGroup.Props>> = ({
     return (
         <Tabs.Root
             className={containerClass}
-            onValueChange={(value) => setSelectedTabIndex(parseInt(value, 10))}
-            defaultValue="0"
+            onValueChange={(value) => {
+                setSelectedTabIndex(value);
+                if (groupId) {
+                    if (selectedGroup) {
+                        const filteredGroupIds = groupIds.filter((group) => group.groupId !== selectedGroup.groupId);
+                        setSelectedTab([{ groupId: groupId, value: value }, ...filteredGroupIds]);
+                    } else {
+                        setSelectedTab([{ groupId: groupId, value: value }, ...groupIds]);
+                    }
+                }
+            }}
+            value={selectedTabIndex}
         >
             <div className="rounded-t-[inherit] bg-tag-default-soft">
                 <div className="mx-px flex min-h-10 items-center justify-between shadow-[inset_0_-1px_0_0] shadow-border-default">

packages/ui/app/src/mdx/plugins/rehypeFernCode.ts

@@ -25,10 +25,16 @@ export function rehypeFernCode(): (tree: Root) => void {
 
             if (isMdxJsxFlowElement(node) && (node.name === "CodeBlocks" || node.name === "CodeGroup")) {
                 const codeBlockItems = visitCodeBlockNodes(node);
+                const attributes = [toAttribute("items", codeBlockItems)];
+
+                if (node.attributes.length > 0 && node.attributes[0]?.value) {
+                    attributes.push(toAttribute("groupId", node.attributes[0]?.value));
+                }
+
                 parent?.children.splice(index, 1, {
                     type: "mdxJsxFlowElement",
                     name: "CodeBlocks",
-                    attributes: [toAttribute("items", codeBlockItems)],
+                    attributes: attributes,
                     children: [],
                 });
                 return "skip";
@@ -107,6 +113,8 @@ export function rehypeFernCode(): (tree: Root) => void {
 
 function visitCodeBlockNodes(nodeToVisit: MdxJsxFlowElementHast) {
     const codeBlockItems: CodeGroup.Item[] = [];
+    let groupId: string | undefined;
+
     visit(nodeToVisit, (node) => {
         if (isMdxJsxFlowElement(node) && node.name === "CodeBlock") {
             const jsxAttributes = node.attributes.filter(
@@ -154,7 +162,6 @@ function visitCodeBlockNodes(nodeToVisit: MdxJsxFlowElementHast) {
                 });
             }
         }
-        return true;
     });
     return codeBlockItems;
 }