Synchronize tabs and allow users to set initial state via query params

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

@@ -0,0 +1,20 @@
+import { useAtomValue } from "jotai";
+import { atomWithStorage } from "jotai/utils";
+
+export type Group = Record<string, string>;
+
+export const FERN_GROUPS = atomWithStorage<Group>("group-id", {});
+FERN_GROUPS.debugLabel = "FERN_GROUPS";
+
+export function useGroup({ key }: { key?: string }): { selectedGroup?: Group; group: Group } {
+    const group = useAtomValue(FERN_GROUPS);
+
+    if (key) {
+        if (key in group && group[key]) {
+            return { selectedGroup: { [key]: group[key] as string }, group };
+        }
+        return { selectedGroup: undefined, group };
+    }
+
+    return { selectedGroup: undefined, group };
+}

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

@@ -3,6 +3,7 @@ export * from "./auth";
 export * from "./docs";
 export * from "./files";
 export * from "./flags";
+export * from "./groups";
 export * from "./hooks";
 export * from "./lang";
 export * from "./layout";

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

@@ -1,10 +1,10 @@
 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 { HorizontalOverflowMask } from "../../../commons/HorizontalOverflowMask";
 import { FernSyntaxHighlighter, FernSyntaxHighlighterProps } from "../../../syntax-highlighting/FernSyntaxHighlighter";
+import { useTabSelection } from "../../hooks";
 
 export declare namespace CodeGroup {
     export interface Item extends FernSyntaxHighlighterProps {
@@ -13,12 +13,13 @@ 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 { selected, setSelected } = useTabSelection({ 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-['']",
@@ -48,11 +49,7 @@ export const CodeGroup: React.FC<React.PropsWithChildren<CodeGroup.Props>> = ({
     }
 
     return (
-        <Tabs.Root
-            className={containerClass}
-            onValueChange={(value) => setSelectedTabIndex(parseInt(value, 10))}
-            defaultValue="0"
-        >
+        <Tabs.Root className={containerClass} onValueChange={setSelected} value={selected}>
             <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">
                     <Tabs.List className="flex min-h-10" asChild>
@@ -71,7 +68,9 @@ export const CodeGroup: React.FC<React.PropsWithChildren<CodeGroup.Props>> = ({
                         </HorizontalOverflowMask>
                     </Tabs.List>
 
-                    <CopyToClipboardButton className="ml-2 mr-1" content={items[selectedTabIndex]?.code} />
+                    <div>
+                        <CopyToClipboardButton className="mx-1" content={items[parseInt(selected)]?.code} />
+                    </div>
                 </div>
             </div>
             {items.map((item, idx) => (

packages/ui/app/src/mdx/components/tabs/Tabs.tsx

@@ -1,7 +1,8 @@
 import * as RadixTabs from "@radix-ui/react-tabs";
 import { slug } from "github-slugger";
 import { useRouter } from "next/router";
-import { FC, ReactNode, useEffect, useState } from "react";
+import { FC, ReactNode, useEffect } from "react";
+import { useTabSelection } from "../../hooks";
 
 export interface TabProps {
     title: string;
@@ -12,23 +13,25 @@ export interface TabProps {
 export interface TabGroupProps {
     tabs: TabProps[];
     toc?: boolean;
+    groupId?: string;
 }
 
-export const TabGroup: FC<TabGroupProps> = ({ tabs, toc: parentToc = true }) => {
-    const [activeTab, setActiveTab] = useState("0");
+export const TabGroup: FC<TabGroupProps> = ({ tabs, groupId, toc: parentToc = true }) => {
     const router = useRouter();
     const anchor = router.asPath.split("#")[1];
+    const { selected, setSelected } = useTabSelection({ groupId });
+
     useEffect(() => {
         if (anchor != null) {
             const anchorTab = tabs.findIndex((tab) => slug(tab.title) === anchor);
             if (anchorTab >= 0) {
-                setActiveTab(anchorTab.toString());
+                setSelected(anchorTab.toString());
             }
         }
-    }, [anchor, tabs]);
+    }, [anchor, tabs, setSelected]);
 
     return (
-        <RadixTabs.Root value={activeTab} onValueChange={setActiveTab}>
+        <RadixTabs.Root value={selected} onValueChange={setSelected}>
             <RadixTabs.List className="border-default mb-6 mt-4 flex gap-4 border-b first:-mt-3">
                 {tabs.map(({ title, toc = parentToc }, idx) => {
                     const id = slug(title);

packages/ui/app/src/mdx/hooks/index.ts

@@ -0,0 +1 @@
+export * from "./useTabSelection";

packages/ui/app/src/mdx/hooks/useTabSelection.ts

@@ -0,0 +1,31 @@
+import { useAtom } from "jotai";
+import { useEffect, useState } from "react";
+import { FERN_GROUPS, Group, useGroup } from "../../atoms";
+
+export const useTabSelection = ({
+    groupId,
+}: {
+    groupId?: string;
+}): { selected: string; setSelected: (val: string) => void } => {
+    const { selectedGroup } = useGroup({ key: groupId });
+    const [_, setGroups] = useAtom<Group>(FERN_GROUPS);
+    const [selectedTabIndex, setSelectedTabIndex] = useState("0");
+
+    useEffect(() => {
+        if (groupId && selectedGroup && selectedGroup[groupId]) {
+            setSelectedTabIndex(selectedGroup[groupId] as string);
+        }
+    }, [groupId, selectedGroup]);
+
+    const setSelected = (value: string) => {
+        if (groupId) {
+            setGroups((prev) => {
+                return { ...prev, [groupId]: value };
+            });
+        } else {
+            setSelectedTabIndex(value);
+        }
+    };
+
+    return { selected: selectedTabIndex, setSelected };
+};

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

@@ -25,10 +25,11 @@ export function rehypeFernCode(): (tree: Root) => void {
 
             if (isMdxJsxFlowElement(node) && (node.name === "CodeBlocks" || node.name === "CodeGroup")) {
                 const codeBlockItems = visitCodeBlockNodes(node);
+
                 parent?.children.splice(index, 1, {
                     type: "mdxJsxFlowElement",
                     name: "CodeBlocks",
-                    attributes: [toAttribute("items", codeBlockItems)],
+                    attributes: [toAttribute("items", codeBlockItems), ...node.attributes],
                     children: [],
                 });
                 return "skip";

packages/ui/components/src/CopyToClipboardButton.tsx

@@ -1,6 +1,7 @@
 import { useCopyToClipboard } from "@fern-ui/react-commons";
 import { CopyIcon } from "@radix-ui/react-icons";
 import cn from "clsx";
+import { ElementType } from "react";
 import { Check } from "react-feather";
 import { FernButton } from "./FernButton";
 import { FernTooltip, FernTooltipProvider } from "./FernTooltip";
@@ -12,6 +13,7 @@ export declare namespace CopyToClipboardButton {
         testId?: string;
         children?: (onClick: ((e: React.MouseEvent) => void) | undefined) => React.ReactNode;
         onClick?: (e: React.MouseEvent) => void;
+        icon?: ElementType;
     }
 }
 
@@ -21,6 +23,7 @@ export const CopyToClipboardButton: React.FC<CopyToClipboardButton.Props> = ({
     testId,
     children,
     onClick,
+    icon: Icon = CopyIcon,
 }) => {
     const { copyToClipboard, wasJustCopied } = useCopyToClipboard(content);
 
@@ -47,7 +50,7 @@ export const CopyToClipboardButton: React.FC<CopyToClipboardButton.Props> = ({
                         }}
                         data-testid={testId}
                         rounded={true}
-                        icon={wasJustCopied ? <Check className="size-4" /> : <CopyIcon className="size-4" />}
+                        icon={wasJustCopied ? <Check className="size-4" /> : <Icon className="size-4" />}
                         variant="minimal"
                         intent={wasJustCopied ? "success" : "none"}
                         disableAutomaticTooltip={true}