From fb5c7660f660dadab4f59aa9d2d98907f2adcb26 Mon Sep 17 00:00:00 2001
From: Denise Li <dli@squareup.com>
Date: Wed, 25 Sep 2024 11:58:04 -0400
Subject: [PATCH] fix: console module tree: rip out disclosure button to fix
 all the state consistency bugs (#2811)

Part 1 of https://github.com/TBD54566975/ftl/issues/2760
---
 frontend/console/e2e/module.spec.ts           |  2 +-
 frontend/console/e2e/verb.spec.ts             |  2 +-
 .../src/features/modules/ModulesTree.tsx      | 77 ++++++++++---------
 3 files changed, 44 insertions(+), 37 deletions(-)

diff --git a/frontend/console/e2e/module.spec.ts b/frontend/console/e2e/module.spec.ts
index 52e9fea8ad..519714efdc 100644
--- a/frontend/console/e2e/module.spec.ts
+++ b/frontend/console/e2e/module.spec.ts
@@ -4,7 +4,7 @@ ftlTest('shows verbs for deployment', async ({ page }) => {
   const modulesNavItem = page.getByRole('link', { name: 'Modules' })
   await modulesNavItem.click()
 
-  const moduleEchoRow = page.getByRole('button', { name: 'echo' })
+  const moduleEchoRow = page.locator('div.cursor-pointer').getByText('echo')
   const moduleEcho = moduleEchoRow.locator('svg').nth(1)
   await moduleEcho.click()
 
diff --git a/frontend/console/e2e/verb.spec.ts b/frontend/console/e2e/verb.spec.ts
index 60e30a7c0f..0b07462015 100644
--- a/frontend/console/e2e/verb.spec.ts
+++ b/frontend/console/e2e/verb.spec.ts
@@ -4,7 +4,7 @@ ftlTest.beforeEach(async ({ page }) => {
   const modulesNavItem = page.getByRole('link', { name: 'Modules' })
   await modulesNavItem.click()
 
-  const moduleEcho = page.getByRole('button', { name: 'echo' })
+  const moduleEcho = page.locator('div.cursor-pointer').getByText('echo').nth(0)
   await moduleEcho.click()
 
   const verbEcho = page.locator('div#decl-echo')
diff --git a/frontend/console/src/features/modules/ModulesTree.tsx b/frontend/console/src/features/modules/ModulesTree.tsx
index 2992d77e9e..020712eee2 100644
--- a/frontend/console/src/features/modules/ModulesTree.tsx
+++ b/frontend/console/src/features/modules/ModulesTree.tsx
@@ -1,6 +1,5 @@
-import { Disclosure, DisclosureButton, DisclosurePanel } from '@headlessui/react'
 import { ArrowRight01Icon, CircleArrowRight02Icon, CodeIcon, FileExportIcon, PackageIcon } from 'hugeicons-react'
-import { useEffect, useMemo, useRef } from 'react'
+import { useEffect, useMemo, useRef, useState } from 'react'
 import { useNavigate, useParams } from 'react-router-dom'
 import type { Decl } from '../../protos/xyz/block/ftl/v1/schema/schema_pb'
 import { classNames } from '../../utils'
@@ -60,7 +59,7 @@ const ModuleSection = ({ module, isExpanded, toggleExpansion }: { module: Module
   const { moduleName, declName } = useParams()
   const navigate = useNavigate()
   const isSelected = useMemo(() => moduleName === module.name, [moduleName, module.name])
-  const moduleRef = useRef<HTMLButtonElement>(null)
+  const moduleRef = useRef<HTMLDivElement>(null)
 
   // Scroll to the selected module on page load
   useEffect(() => {
@@ -75,54 +74,62 @@ const ModuleSection = ({ module, isExpanded, toggleExpansion }: { module: Module
 
   return (
     <li key={module.name} id={`module-tree-module-${module.name}`} className='my-2'>
-      <Disclosure as='div' defaultOpen={isExpanded}>
-        <DisclosureButton
-          ref={moduleRef}
-          className={classNames(
-            isSelected ? 'bg-gray-200 dark:bg-gray-700 hover:bg-gray-300 hover:dark:bg-gray-600' : 'hover:bg-gray-200 hover:dark:bg-gray-700',
-            'group flex w-full modules-center gap-x-2 space-y-1 text-left text-sm font-medium leading-6',
-          )}
-          onClick={() => toggleExpansion(module.name)}
-        >
-          <PackageIcon aria-hidden='true' className='size-4 my-1 ml-3 shrink-0' />
-          {module.name}
-          <CircleArrowRight02Icon
-            className='size-4 shrink-0 rounded-md hover:bg-gray-200 dark:hover:bg-gray-600'
-            onClick={(e) => {
-              e.preventDefault()
-              e.stopPropagation()
-              navigate(`/modules/${module.name}`)
-            }}
-          />
-          {module.decls.length === 0 || (
-            <ArrowRight01Icon aria-hidden='true' className='ml-auto mr-2 h-4 w-4 shrink-0 group-data-[open]:rotate-90 group-data-[open]:text-gray-500' />
-          )}
-        </DisclosureButton>
-        <DisclosurePanel as='ul'>
+      <div
+        ref={moduleRef}
+        className={classNames(
+          isSelected ? 'bg-gray-200 dark:bg-gray-700 hover:bg-gray-300 hover:dark:bg-gray-600' : 'hover:bg-gray-200 hover:dark:bg-gray-700',
+          'group flex w-full modules-center gap-x-2 space-y-1 text-left text-sm font-medium cursor-pointer leading-6',
+        )}
+        onClick={() => toggleExpansion(module.name)}
+      >
+        <PackageIcon aria-hidden='true' className='size-4 my-1 ml-3 shrink-0' />
+        {module.name}
+        <CircleArrowRight02Icon
+          className='size-4 shrink-0 rounded-md hover:bg-gray-200 dark:hover:bg-gray-600'
+          onClick={(e) => {
+            e.preventDefault()
+            e.stopPropagation()
+            navigate(`/modules/${module.name}`)
+          }}
+        />
+        {module.decls.length === 0 || (
+          <ArrowRight01Icon aria-hidden='true' className='ml-auto mr-2 h-4 w-4 shrink-0 group-data-[open]:rotate-90 group-data-[open]:text-gray-500' />
+        )}
+      </div>
+      {isExpanded && (
+        <ul>
           {module.decls.map((d, i) => (
             <DeclNode key={i} decl={d} href={declUrl(module.name, d)} isSelected={isSelected && declName === d.value.value?.name} />
           ))}
-        </DisclosurePanel>
-      </Disclosure>
+        </ul>
+      )}
     </li>
   )
 }
 
 export const ModulesTree = ({ modules }: { modules: ModuleTreeItem[] }) => {
-  const { moduleName } = useParams()
+  const { moduleName, declName } = useParams()
+
+  const [expandedModules, setExpandedModules] = useState(listExpandedModulesFromLocalStorage())
   useEffect(() => {
-    addModuleToLocalStorageIfMissing(moduleName)
-  }, [moduleName])
+    if (declName) {
+      addModuleToLocalStorageIfMissing(moduleName)
+    }
+  }, [moduleName, declName])
 
-  modules.sort((m1, m2) => Number(m1.isBuiltin) - Number(m2.isBuiltin))
+  function toggleFn(moduleName: string) {
+    toggleModuleExpansionInLocalStorage(moduleName)
+    setExpandedModules(listExpandedModulesFromLocalStorage())
+    return
+  }
 
-  const expandedModules = listExpandedModulesFromLocalStorage()
+  modules.sort((m1, m2) => Number(m1.isBuiltin) - Number(m2.isBuiltin))
   return (
     <div className={'flex grow flex-col h-full gap-y-5 overflow-y-auto bg-gray-100 dark:bg-gray-900'}>
       <nav>
         <ul>
           {modules.map((m) => (
-            <ModuleSection key={m.name} module={m} isExpanded={expandedModules.includes(m.name)} toggleExpansion={toggleModuleExpansionInLocalStorage} />
+            <ModuleSection key={m.name} module={m} isExpanded={expandedModules.includes(m.name)} toggleExpansion={toggleFn} />
           ))}
         </ul>
       </nav>