Got enhancing working in desktop2

.github/actions/pnpm_install/action.yaml

@@ -3,7 +3,7 @@ runs:
   steps:
     - uses: pnpm/action-setup@v4
       with:
-        version: "10.18.3"
+        version: "10.19.0"
         run_install: false
 
     - uses: actions/setup-node@v4

apps/desktop/package.json

@@ -36,6 +36,7 @@
     "@hypr/utils": "workspace:^",
     "@iconify-icon/react": "^3.0.1",
     "@lobehub/icons": "^2.43.1",
+    "@openrouter/ai-sdk-provider": "^1.2.0",
     "@orama/highlight": "^0.1.9",
     "@orama/orama": "^3.1.16",
     "@orama/plugin-qps": "^3.1.16",

apps/desktop/src/chat/transport.ts

@@ -1,31 +1,19 @@
-import type { ChatRequestOptions, ChatTransport, LanguageModel, UIMessageChunk } from "ai";
-import { convertToModelMessages, smoothStream, stepCountIs, streamText } from "ai";
+import type { ChatTransport, LanguageModel } from "ai";
+import { convertToModelMessages, Experimental_Agent as Agent, stepCountIs } from "ai";
 
 import { ToolRegistry } from "../contexts/tool";
 import type { HyprUIMessage } from "./types";
 
 export class CustomChatTransport implements ChatTransport<HyprUIMessage> {
   constructor(private registry: ToolRegistry, private model: LanguageModel) {}
 
-  async sendMessages(
-    options:
-      & {
-        chatId: string;
-        messages: HyprUIMessage[];
-        abortSignal: AbortSignal | undefined;
-      }
-      & { trigger: "submit-message" | "regenerate-message"; messageId: string | undefined }
-      & ChatRequestOptions,
-  ): Promise<ReadableStream<UIMessageChunk>> {
+  sendMessages: ChatTransport<HyprUIMessage>["sendMessages"] = async (options) => {
     const tools = this.registry.getForTransport();
 
-    const result = streamText({
+    const agent = new Agent({
       model: this.model,
-      messages: convertToModelMessages(options.messages),
-      experimental_transform: smoothStream({ chunking: "word" }),
       tools,
       stopWhen: stepCountIs(5),
-      abortSignal: options.abortSignal,
       prepareStep: async ({ messages }) => {
         if (messages.length > 20) {
           return { messages: messages.slice(-10) };
@@ -35,6 +23,8 @@ export class CustomChatTransport implements ChatTransport<HyprUIMessage> {
       },
     });
 
+    const result = agent.stream({ messages: convertToModelMessages(options.messages) });
+
     return result.toUIMessageStream({
       originalMessages: options.messages,
       messageMetadata: ({ part }) => {
@@ -47,9 +37,9 @@ export class CustomChatTransport implements ChatTransport<HyprUIMessage> {
         return error instanceof Error ? error.message : String(error);
       },
     });
-  }
+  };
 
-  async reconnectToStream(): Promise<ReadableStream<UIMessageChunk> | null> {
+  reconnectToStream: ChatTransport<HyprUIMessage>["reconnectToStream"] = async () => {
     return null;
-  }
+  };
 }

apps/desktop/src/components/main/body/index.tsx

@@ -270,12 +270,21 @@ function TabChatButton() {
 }
 
 export function StandardTabWrapper(
-  { children, afterBorder }: { children: React.ReactNode; afterBorder?: React.ReactNode },
+  {
+    children,
+    afterBorder,
+    floatingButton,
+  }: {
+    children: React.ReactNode;
+    afterBorder?: React.ReactNode;
+    floatingButton?: React.ReactNode;
+  },
 ) {
   return (
     <div className="flex flex-col h-full">
       <div className="flex flex-col rounded-lg border flex-1 overflow-hidden relative">
         {children}
+        {floatingButton}
         <TabChatButton />
       </div>
       {afterBorder}

apps/desktop/src/components/main/body/sessions/floating/generate.tsx

@@ -1,20 +1,54 @@
+import { cn } from "@hypr/utils";
 import { SparklesIcon } from "lucide-react";
 import { useState } from "react";
 
-import { cn } from "@hypr/utils";
+import { useAITask } from "../../../../../contexts/ai-task";
+import { useLanguageModel } from "../../../../../hooks/useLLMConnection";
 import * as persisted from "../../../../../store/tinybase/persisted";
 
 import { FloatingButton } from "./shared";
 
-export function GenerateButton() {
+export function GenerateButton({ sessionId }: { sessionId: string }) {
   const [showTemplates, setShowTemplates] = useState(false);
+  const model = useLanguageModel();
+
+  const taskId = `${sessionId}-enhance`;
+
+  const { generate, status } = useAITask((state) => ({
+    generate: state.generate,
+    status: state.tasks[taskId]?.status ?? "idle",
+  }));
 
   const templates = persisted.UI.useResultTable(persisted.QUERIES.visibleTemplates, persisted.STORE_ID);
+  const rawMd = persisted.UI.useCell("sessions", sessionId, "raw_md", persisted.STORE_ID);
+
+  const updateEnhancedMd = persisted.UI.useSetPartialRowCallback(
+    "sessions",
+    sessionId,
+    (input: string) => ({ enhanced_md: input }),
+    [],
+    persisted.STORE_ID,
+  );
+
+  const onRegenerate = async (_templateId: string | null) => {
+    if (!model) {
+      return;
+    }
 
-  const onRegenerate = (templateId: string | null) => {
-    console.log("Regenerate clicked:", templateId);
+    await generate(taskId, {
+      model,
+      taskType: "enhance",
+      args: { rawMd },
+      onComplete: updateEnhancedMd,
+    });
   };
 
+  const isGenerating = status === "generating";
+
+  if (isGenerating) {
+    return null;
+  }
+
   return (
     <div>
       <div
@@ -54,17 +88,19 @@ export function GenerateButton() {
         </div>
       </div>
 
-      <FloatingButton
-        icon={<SparklesIcon className="w-4 h-4" />}
-        onMouseEnter={() => setShowTemplates(true)}
-        onMouseLeave={() => setShowTemplates(false)}
-        onClick={() => {
-          setShowTemplates(false);
-          onRegenerate(null);
-        }}
-      >
-        Regenerate
-      </FloatingButton>
+      <div className="flex flex-col items-center">
+        <FloatingButton
+          icon={<SparklesIcon className="w-4 h-4" />}
+          onMouseEnter={() => setShowTemplates(true)}
+          onMouseLeave={() => setShowTemplates(false)}
+          onClick={() => {
+            setShowTemplates(false);
+            onRegenerate(null);
+          }}
+        >
+          <span>Regenerate</span>
+        </FloatingButton>
+      </div>
     </div>
   );
 }

apps/desktop/src/components/main/body/sessions/floating/index.tsx

@@ -19,7 +19,7 @@ export function FloatingActionButton({ tab }: { tab: Extract<Tab, { type: "sessi
   } else if (tab.state.editor === "enhanced") {
     return (
       <FloatingButtonContainer>
-        <GenerateButton />
+        <GenerateButton sessionId={tab.id} />
       </FloatingButtonContainer>
     );
   } else if (tab.state.editor === "transcript") {

apps/desktop/src/components/main/body/sessions/index.tsx

@@ -53,15 +53,17 @@ export function TabContentNote({ tab }: { tab: Extract<Tab, { type: "sessions" }
 
   return (
     <AudioPlayer.Provider url={audioUrl ?? ""}>
-      <StandardTabWrapper afterBorder={tab.state.editor === "transcript" && <AudioPlayer.Timeline />}>
-        <div className="p-2">
+      <StandardTabWrapper
+        afterBorder={tab.state.editor === "transcript" && <AudioPlayer.Timeline />}
+        floatingButton={<FloatingActionButton tab={tab} />}
+      >
+        <div className="flex flex-col h-full p-2">
           <OuterHeader sessionId={tab.id} />
-          <div className="mt-3 px-2">
+          <div className="mt-3 px-2 flex-shrink-0">
             <TitleInput tab={tab} />
-            <div className="mt-2">
-              <NoteInput tab={tab} />
-            </div>
-            <FloatingActionButton tab={tab} />
+          </div>
+          <div className="mt-2 px-2 flex-1 min-h-0">
+            <NoteInput tab={tab} />
           </div>
         </div>
       </StandardTabWrapper>

apps/desktop/src/components/main/body/sessions/note-input/enhanced/editor.tsx

@@ -0,0 +1,36 @@
+import { forwardRef } from "react";
+
+import { TiptapEditor } from "@hypr/tiptap/editor";
+import NoteEditor from "@hypr/tiptap/editor";
+import * as persisted from "../../../../../../store/tinybase/persisted";
+
+export const EnhancedEditor = forwardRef<{ editor: TiptapEditor | null }, { sessionId: string }>(
+  ({ sessionId }, ref) => {
+    const value = persisted.UI.useCell("sessions", sessionId, "enhanced_md", persisted.STORE_ID);
+
+    const handleEnhancedChange = persisted.UI.useSetPartialRowCallback(
+      "sessions",
+      sessionId,
+      (input: string) => ({ enhanced_md: input }),
+      [],
+      persisted.STORE_ID,
+    );
+
+    return (
+      <div className="h-full">
+        <NoteEditor
+          ref={ref}
+          key={`session-${sessionId}-enhanced`}
+          initialContent={value ?? ""}
+          handleChange={handleEnhancedChange}
+          mentionConfig={{
+            trigger: "@",
+            handleSearch: async () => {
+              return [];
+            },
+          }}
+        />
+      </div>
+    );
+  },
+);

apps/desktop/src/components/main/body/sessions/note-input/enhanced/index.tsx

@@ -0,0 +1,28 @@
+import { type TiptapEditor } from "@hypr/tiptap/editor";
+import { forwardRef } from "react";
+
+import { useAITask } from "../../../../../../contexts/ai-task";
+import { EnhancedEditor } from "./editor";
+import { StreamingView } from "./streaming";
+
+export const Enhanced = forwardRef<
+  { editor: TiptapEditor | null },
+  { sessionId: string }
+>(({ sessionId }, ref) => {
+  const taskId = `${sessionId}-enhance`;
+
+  const { status, error } = useAITask((state) => ({
+    status: state.tasks[taskId]?.status ?? "idle",
+    error: state.tasks[taskId]?.error,
+  }));
+
+  if (status === "error" && error) {
+    return <pre>{error.message}</pre>;
+  }
+
+  if (status === "generating") {
+    return <StreamingView sessionId={sessionId} />;
+  }
+
+  return <EnhancedEditor ref={ref} sessionId={sessionId} />;
+});