From 877926bc4fd40e7193629f45bb3632d8ce08a044 Mon Sep 17 00:00:00 2001
From: Aaron Villalpando <aaron@boundaryml.com>
Date: Tue, 2 Jul 2024 17:16:40 -0700
Subject: [PATCH] fix initial selected file

---
 .../app/[project_id]/_atoms/atoms.ts          |  6 ------
 .../_components/CodeMirrorEditor.tsx          | 19 ++++++++++++++++++-
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/typescript/fiddle-frontend/app/[project_id]/_atoms/atoms.ts b/typescript/fiddle-frontend/app/[project_id]/_atoms/atoms.ts
index 3c9aef10c..0ca5e5765 100644
--- a/typescript/fiddle-frontend/app/[project_id]/_atoms/atoms.ts
+++ b/typescript/fiddle-frontend/app/[project_id]/_atoms/atoms.ts
@@ -38,12 +38,6 @@ export const activeFileNameAtom = atom(
     const files = get(currentEditorFilesAtom)
     if (files.some((f) => f.path === path)) {
       set(activeFileNameAtomRaw, path)
-      const availableFunctions = get(availableFunctionsAtom)
-      // Set the default selected function when the current active file is changed
-      const selectedFunction = availableFunctions.find((f) => f.span.file_path === path)
-      if (selectedFunction) {
-        set(selectedFunctionAtom, selectedFunction.name)
-      }
     }
   },
 )
diff --git a/typescript/fiddle-frontend/app/[project_id]/_components/CodeMirrorEditor.tsx b/typescript/fiddle-frontend/app/[project_id]/_components/CodeMirrorEditor.tsx
index 780a295d2..6d755e100 100644
--- a/typescript/fiddle-frontend/app/[project_id]/_components/CodeMirrorEditor.tsx
+++ b/typescript/fiddle-frontend/app/[project_id]/_components/CodeMirrorEditor.tsx
@@ -4,7 +4,13 @@ import { BAML_DIR } from '@/lib/constants'
 import type { BAMLProject } from '@/lib/exampleProjects'
 import { BAML, theme } from '@baml/codemirror-lang'
 import type { ParserDatabase } from '@baml/common'
-import { diagnositicsAtom, numErrorsAtom, updateFileAtom } from '@baml/playground-common/baml_wasm_web/EventListener'
+import {
+  availableFunctionsAtom,
+  diagnositicsAtom,
+  numErrorsAtom,
+  selectedFunctionAtom,
+  updateFileAtom,
+} from '@baml/playground-common/baml_wasm_web/EventListener'
 import { atomStore } from '@baml/playground-common/baml_wasm_web/JotaiProvider'
 import { projectFamilyAtom, runtimeFamilyAtom } from '@baml/playground-common/baml_wasm_web/baseAtoms'
 import { Button } from '@baml/playground-common/components/ui/button'
@@ -142,6 +148,9 @@ export const CodeMirrorEditor = ({ project }: { project: BAMLProject }) => {
   const activeFileContent = useAtomValue(activeFileContentAtom)
   const updateFile = useSetAtom(updateFileAtom)
 
+  const availableFunctions = useAtomValue(availableFunctionsAtom)
+  const setSelectedFunction = useSetAtom(selectedFunctionAtom)
+
   const ref = useRef<ReactCodeMirrorRef>({})
 
   // force linting on file changes so playground updates
@@ -154,6 +163,14 @@ export const CodeMirrorEditor = ({ project }: { project: BAMLProject }) => {
     }
   }, [JSON.stringify(editorFiles)])
 
+  useEffect(() => {
+    const func = availableFunctions.find((f) => f.span.file_path === activeFile)
+    if (func) {
+      console.log('setting selected function', func.name)
+      setSelectedFunction(func.name)
+    }
+  }, [JSON.stringify(editorFiles.map((f) => f.path)), activeFile, availableFunctions])
+
   const setUnsavedChanges = useSetAtom(unsavedChangesAtom)
 
   const langExtensions = getLanguage(activeFile)