Template after recording + Refactored Enhancement Codes

apps/desktop/src/components/editor-area/floating-button.tsx

@@ -147,7 +147,6 @@ export function FloatingButton({
 
       const handleWindowFocus = () => {
         queryClient.invalidateQueries({ queryKey: ["templates"] });
-        queryClient.invalidateQueries({ queryKey: ["llm-connection"] });
         window.removeEventListener("focus", handleWindowFocus);
       };
 

apps/desktop/src/components/editor-area/index.tsx

@@ -63,37 +63,12 @@ export default function EditorArea({
     [sessionId, showRaw],
   );
 
-  const [needsRestoration, setNeedsRestoration] = useState(false);
-  const [originalTemplateId, setOriginalTemplateId] = useState<string | null>(null);
-  const queryClient = useQueryClient();
-
-  const configQuery = useQuery({
-    queryKey: ["config", "general"],
-    queryFn: async () => {
-      const result = await dbCommands.getConfig();
-      return result;
-    },
-  });
-
-  const setConfigMutation = useMutation({
-    mutationFn: async (configData: any) => {
-      await dbCommands.setConfig(configData);
-    },
-    onSuccess: () => {
-      queryClient.invalidateQueries({ queryKey: ["config", "general"] });
-    },
-    onError: (error) => {
-      console.error("Failed to set template config:", error);
-    },
-  });
-
   const templatesQuery = useQuery({
     queryKey: ["templates"],
     queryFn: () => dbCommands.listTemplates(),
     refetchOnWindowFocus: true,
   });
 
-  const generateTitle = useGenerateTitleMutation({ sessionId });
   const preMeetingNote = useSession(sessionId, (s) => s.session.pre_meeting_memo_html) ?? "";
   const hasTranscriptWords = useSession(sessionId, (s) => s.session.words.length > 0);
 
@@ -108,27 +83,14 @@ export default function EditorArea({
     rawContent,
     isLocalLlm: llmConnectionQuery.data?.type === "HyprLocal",
     onSuccess: (content) => {
-      generateTitle.mutate({ enhancedContent: content });
-
-      if (needsRestoration && configQuery.data) {
-        const restoreConfig = {
-          ...configQuery.data,
-          general: {
-            ...configQuery.data.general,
-            selected_template_id: originalTemplateId,
-          },
-        };
-        setConfigMutation.mutate(restoreConfig);
-        setNeedsRestoration(false);
-        setOriginalTemplateId(null);
-      }
-
       if (hasTranscriptWords) {
         generateTitle.mutate({ enhancedContent: content });
       }
     },
   });
 
+  const generateTitle = useGenerateTitleMutation({ sessionId });
+
   useAutoEnhance({
     sessionId,
     enhanceStatus: enhance.status,
@@ -151,46 +113,13 @@ export default function EditorArea({
     [showRaw, enhancedContent, rawContent],
   );
 
-  const handleEnhanceWithTemplate = useCallback(async (templateId: string) => {
-    if (configQuery.data) {
-      const currentTemplateId = configQuery.data.general?.selected_template_id || null;
-      setOriginalTemplateId(currentTemplateId);
-      setNeedsRestoration(true);
-
-      const targetTemplateId = templateId === "auto" ? null : templateId;
-
-      const updatedConfig = {
-        ...configQuery.data,
-        general: {
-          ...configQuery.data.general,
-          selected_template_id: targetTemplateId,
-        },
-      };
-
-      try {
-        await setConfigMutation.mutateAsync(updatedConfig);
-
-        await new Promise(resolve => setTimeout(resolve, 200));
-
-        const verifyConfig = await dbCommands.getConfig();
-
-        if (verifyConfig.general?.selected_template_id !== targetTemplateId) {
-          setOriginalTemplateId(null);
-          setNeedsRestoration(false);
-          return;
-        }
-      } catch (error) {
-        setOriginalTemplateId(null);
-        setNeedsRestoration(false);
-        return;
-      }
-    }
-
-    enhance.mutate();
-  }, [enhance, configQuery.data, setConfigMutation]);
+  const handleEnhanceWithTemplate = useCallback((templateId: string) => {
+    const targetTemplateId = templateId === "auto" ? null : templateId;
+    enhance.mutate({ templateId: targetTemplateId, triggerType: "template" });
+  }, [enhance]);
 
   const handleClickEnhance = useCallback(() => {
-    enhance.mutate();
+    enhance.mutate({ triggerType: "manual" });
   }, [enhance]);
 
   const safelyFocusEditor = useCallback(() => {
@@ -317,7 +246,13 @@ export function useEnhanceMutation({
 
   const enhance = useMutation({
     mutationKey: ["enhance", sessionId],
-    mutationFn: async () => {
+    mutationFn: async ({
+      triggerType,
+      templateId,
+    }: {
+      triggerType: "manual" | "template" | "auto";
+      templateId?: string | null;
+    } = { triggerType: "manual" }) => {
       await queryClient.invalidateQueries({ queryKey: ["llm-connection"] });
       await new Promise(resolve => setTimeout(resolve, 100));
 
@@ -347,15 +282,20 @@ export function useEnhanceMutation({
         return;
       }
 
+      // Get current config for default template
       const config = await dbCommands.getConfig();
 
+      // Use provided templateId or fall back to config
+      const effectiveTemplateId = templateId !== undefined
+        ? templateId
+        : config.general?.selected_template_id;
+
       let templateInfo = "";
       let customGrammar: string | null = null;
-      const selectedTemplateId = config.general.selected_template_id;
 
-      if (selectedTemplateId) {
+      if (effectiveTemplateId) {
         const templates = await dbCommands.listTemplates();
-        const selectedTemplate = templates.find(t => t.id === selectedTemplateId);
+        const selectedTemplate = templates.find(t => t.id === effectiveTemplateId);
 
         if (selectedTemplate) {
           if (selectedTemplate.sections && selectedTemplate.sections.length > 0) {
@@ -560,9 +500,11 @@ function useAutoEnhance({
 }: {
   sessionId: string;
   enhanceStatus: string;
-  enhanceMutate: () => void;
+  enhanceMutate: (params: { triggerType: "auto"; templateId?: string | null }) => void;
 }) {
   const ongoingSessionStatus = useOngoingSession((s) => s.status);
+  const autoEnhanceTemplate = useOngoingSession((s) => s.autoEnhanceTemplate);
+  const setAutoEnhanceTemplate = useOngoingSession((s) => s.setAutoEnhanceTemplate);
   const prevOngoingSessionStatus = usePreviousValue(ongoingSessionStatus);
   const setShowRaw = useSession(sessionId, (s) => s.setShowRaw);
 
@@ -573,14 +515,25 @@ function useAutoEnhance({
       && enhanceStatus !== "pending"
     ) {
       setShowRaw(false);
-      enhanceMutate();
+
+      // Use the selected template and then clear it
+      enhanceMutate({
+        triggerType: "auto",
+        templateId: autoEnhanceTemplate,
+      });
+
+      // Clear the template after using it (one-time use)
+      setAutoEnhanceTemplate(null);
     }
   }, [
     ongoingSessionStatus,
     enhanceStatus,
     sessionId,
     enhanceMutate,
     setShowRaw,
+    autoEnhanceTemplate,
+    setAutoEnhanceTemplate,
+    prevOngoingSessionStatus,
   ]);
 }
 

apps/desktop/src/components/editor-area/note-header/listen-button.tsx

@@ -6,10 +6,12 @@ import { useEffect, useState } from "react";
 import SoundIndicator from "@/components/sound-indicator";
 import { useHypr } from "@/contexts";
 import { useEnhancePendingState } from "@/hooks/enhance-pending";
+import { commands as dbCommands } from "@hypr/plugin-db";
 import { commands as listenerCommands } from "@hypr/plugin-listener";
 import { commands as localSttCommands } from "@hypr/plugin-local-stt";
 import { Button } from "@hypr/ui/components/ui/button";
 import { Popover, PopoverContent, PopoverTrigger } from "@hypr/ui/components/ui/popover";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@hypr/ui/components/ui/select";
 import { Spinner } from "@hypr/ui/components/ui/spinner";
 import { sonnerToast, toast } from "@hypr/ui/components/ui/toast";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@hypr/ui/components/ui/tooltip";
@@ -222,6 +224,7 @@ function WhenActive() {
   const ongoingSessionStore = useOngoingSession((s) => ({
     pause: s.pause,
     stop: s.stop,
+    setAutoEnhanceTemplate: s.setAutoEnhanceTemplate,
   }));
   const sessionWords = useSession(ongoingSessionId!, (s) => s.session.words);
   const [isPopoverOpen, setIsPopoverOpen] = useState(false);
@@ -231,7 +234,11 @@ function WhenActive() {
     setIsPopoverOpen(false);
   };
 
-  const handleStopSession = () => {
+  const handleStopSession = (templateId?: string | null) => {
+    if (templateId !== undefined) {
+      ongoingSessionStore.setAutoEnhanceTemplate(templateId);
+    }
+
     ongoingSessionStore.stop();
     setIsPopoverOpen(false);
 
@@ -241,7 +248,7 @@ function WhenActive() {
   };
 
   return (
-    <Popover>
+    <Popover open={isPopoverOpen} onOpenChange={setIsPopoverOpen}>
       <PopoverTrigger asChild>
         <button
           className={cn([
@@ -268,12 +275,13 @@ function RecordingControls({
   onStop,
 }: {
   onPause: () => void;
-  onStop: () => void;
+  onStop: (templateId?: string | null) => void;
 }) {
   const ongoingSessionMuted = useOngoingSession((s) => ({
     micMuted: s.micMuted,
     speakerMuted: s.speakerMuted,
   }));
+  const [selectedTemplate, setSelectedTemplate] = useState<string>("auto");
 
   const toggleMicMuted = useMutation({
     mutationFn: () => listenerCommands.setMicMuted(!ongoingSessionMuted.micMuted),
@@ -283,9 +291,34 @@ function RecordingControls({
     mutationFn: () => listenerCommands.setSpeakerMuted(!ongoingSessionMuted.speakerMuted),
   });
 
+  const configQuery = useQuery({
+    queryKey: ["config"],
+    queryFn: () => dbCommands.getConfig(),
+    refetchOnWindowFocus: true,
+  });
+
+  const templatesQuery = useQuery({
+    queryKey: ["templates"],
+    queryFn: () => dbCommands.listTemplates(),
+    refetchOnWindowFocus: true,
+  });
+
+  useEffect(() => {
+    if (configQuery.data?.general?.selected_template_id) {
+      setSelectedTemplate(configQuery.data.general.selected_template_id);
+    } else {
+      setSelectedTemplate("auto");
+    }
+  }, [configQuery.data]);
+
+  const handleStopWithTemplate = () => {
+    const actualTemplateId = selectedTemplate === "auto" ? null : selectedTemplate;
+    onStop(actualTemplateId);
+  };
+
   return (
     <>
-      <div className="flex w-full justify-between mb-4">
+      <div className="flex w-full justify-between mb-3">
         <AudioControlButton
           isMuted={ongoingSessionMuted.micMuted}
           onClick={() => toggleMicMuted.mutate()}
@@ -298,6 +331,24 @@ function RecordingControls({
         />
       </div>
 
+      <div className="mb-3">
+        <Select value={selectedTemplate} onValueChange={setSelectedTemplate}>
+          <SelectTrigger className="w-full text-sm">
+            <SelectValue placeholder="Select template..." />
+          </SelectTrigger>
+          <SelectContent>
+            <SelectItem value="auto">
+              <Trans>No Template (Default)</Trans>
+            </SelectItem>
+            {templatesQuery.data?.map((template) => (
+              <SelectItem key={template.id} value={template.id}>
+                {template.title || "Untitled"}
+              </SelectItem>
+            ))}
+          </SelectContent>
+        </Select>
+      </div>
+
       <div className="flex gap-2">
         <Button
           variant="outline"
@@ -309,7 +360,7 @@ function RecordingControls({
         </Button>
         <Button
           variant="destructive"
-          onClick={onStop}
+          onClick={handleStopWithTemplate}
           className="w-full"
         >
           <StopCircleIcon size={16} />