From b2e3631a295f810d6f66523a6c84dabf51d9c285 Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Fri, 18 Apr 2025 12:32:57 -0700
Subject: [PATCH] Refactor chat input & properly update state for proposal

---
 src/atoms/chatAtoms.ts                   |  1 +
 src/atoms/proposalAtoms.ts               |  4 ++
 src/components/chat/ChatInput.tsx        | 35 ++++++----
 src/components/chat/HomeChatInput.tsx    | 85 ++++++++++++++++++++++++
 src/hooks/useProposal.ts                 | 28 ++++----
 src/hooks/useStreamChat.ts               |  3 +-
 src/ipc/handlers/chat_stream_handlers.ts | 10 +--
 src/ipc/ipc_client.ts                    |  8 +--
 src/lib/schemas.ts                       |  5 ++
 src/pages/home.tsx                       |  9 ++-
 10 files changed, 140 insertions(+), 48 deletions(-)
 create mode 100644 src/atoms/proposalAtoms.ts
 create mode 100644 src/components/chat/HomeChatInput.tsx

diff --git a/src/atoms/chatAtoms.ts b/src/atoms/chatAtoms.ts
index 6a26399..36aac80 100644
--- a/src/atoms/chatAtoms.ts
+++ b/src/atoms/chatAtoms.ts
@@ -11,6 +11,7 @@ export const selectedChatIdAtom = atom<number | null>(null);
 
 export const isStreamingAtom = atom<boolean>(false);
 export const chatInputValueAtom = atom<string>("");
+export const homeChatInputValueAtom = atom<string>("");
 
 // Atoms for chat list management
 export const chatsAtom = atom<ChatSummary[]>([]);
diff --git a/src/atoms/proposalAtoms.ts b/src/atoms/proposalAtoms.ts
new file mode 100644
index 0000000..d9f19ae
--- /dev/null
+++ b/src/atoms/proposalAtoms.ts
@@ -0,0 +1,4 @@
+import { atom } from "jotai";
+import type { Proposal, ProposalResult } from "@/lib/schemas";
+
+export const proposalResultAtom = atom<ProposalResult | null>(null);
diff --git a/src/components/chat/ChatInput.tsx b/src/components/chat/ChatInput.tsx
index 8c24a52..f347c71 100644
--- a/src/components/chat/ChatInput.tsx
+++ b/src/components/chat/ChatInput.tsx
@@ -11,12 +11,12 @@ import {
   Loader2,
 } from "lucide-react";
 import type React from "react";
-import { useEffect, useRef, useState } from "react";
+import { useCallback, useEffect, useRef, useState } from "react";
 import { ModelPicker } from "@/components/ModelPicker";
 import { useSettings } from "@/hooks/useSettings";
 import { IpcClient } from "@/ipc/ipc_client";
-import { chatInputValueAtom } from "@/atoms/chatAtoms";
-import { useAtom } from "jotai";
+import { chatInputValueAtom, chatMessagesAtom } from "@/atoms/chatAtoms";
+import { useAtom, useSetAtom } from "jotai";
 import { useStreamChat } from "@/hooks/useStreamChat";
 import { useChats } from "@/hooks/useChats";
 import { selectedAppIdAtom } from "@/atoms/appAtoms";
@@ -26,7 +26,8 @@ import { Label } from "@/components/ui/label";
 import { Switch } from "@/components/ui/switch";
 import { useProposal } from "@/hooks/useProposal";
 import { Proposal } from "@/lib/schemas";
-
+import type { Message } from "@/ipc/ipc_types";
+import { isPreviewOpenAtom } from "@/atoms/viewAtoms";
 interface ChatInputActionsProps {
   proposal: Proposal;
   onApprove: () => void;
@@ -36,12 +37,7 @@ interface ChatInputActionsProps {
   isRejecting: boolean; // State for rejecting
 }
 
-interface ChatInputProps {
-  chatId?: number;
-  onSubmit?: () => void;
-}
-
-export function ChatInput({ chatId, onSubmit }: ChatInputProps) {
+export function ChatInput({ chatId }: { chatId?: number }) {
   const [inputValue, setInputValue] = useAtom(chatInputValueAtom);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const { settings, updateSettings, isAnyProviderSetup } = useSettings();
@@ -51,6 +47,8 @@ export function ChatInput({ chatId, onSubmit }: ChatInputProps) {
   const [showError, setShowError] = useState(true);
   const [isApproving, setIsApproving] = useState(false); // State for approving
   const [isRejecting, setIsRejecting] = useState(false); // State for rejecting
+  const [messages, setMessages] = useAtom<Message[]>(chatMessagesAtom);
+  const setIsPreviewOpen = useSetAtom(isPreviewOpenAtom);
 
   // Use the hook to fetch the proposal
   const {
@@ -80,10 +78,19 @@ export function ChatInput({ chatId, onSubmit }: ChatInputProps) {
     }
   }, [error]);
 
+  const fetchChatMessages = useCallback(async () => {
+    if (!chatId) {
+      setMessages([]);
+      return;
+    }
+    const chat = await IpcClient.getInstance().getChat(chatId);
+    setMessages(chat.messages);
+  }, [chatId, setMessages]);
+
   const handleKeyPress = (e: React.KeyboardEvent) => {
     if (e.key === "Enter" && !e.shiftKey) {
       e.preventDefault();
-      submitHandler();
+      handleSubmit();
     }
   };
 
@@ -96,7 +103,6 @@ export function ChatInput({ chatId, onSubmit }: ChatInputProps) {
     setInputValue("");
     await streamMessage({ prompt: currentInput, chatId });
   };
-  const submitHandler = onSubmit ? onSubmit : handleSubmit;
 
   const handleCancel = () => {
     if (chatId) {
@@ -133,7 +139,9 @@ export function ChatInput({ chatId, onSubmit }: ChatInputProps) {
       setError((err as Error)?.message || "An error occurred while approving");
     } finally {
       setIsApproving(false);
+      setIsPreviewOpen(true);
       refreshProposal();
+      fetchChatMessages();
     }
   };
 
@@ -162,6 +170,7 @@ export function ChatInput({ chatId, onSubmit }: ChatInputProps) {
     } finally {
       setIsRejecting(false);
       refreshProposal();
+      fetchChatMessages();
     }
   };
 
@@ -236,7 +245,7 @@ export function ChatInput({ chatId, onSubmit }: ChatInputProps) {
               </button>
             ) : (
               <button
-                onClick={submitHandler}
+                onClick={handleSubmit}
                 disabled={!inputValue.trim() || !isAnyProviderSetup()}
                 className="px-2 py-2 mt-1 mr-2 hover:bg-(--background-darkest) text-(--sidebar-accent-fg) rounded-lg disabled:opacity-50"
               >
diff --git a/src/components/chat/HomeChatInput.tsx b/src/components/chat/HomeChatInput.tsx
new file mode 100644
index 0000000..3eabe3c
--- /dev/null
+++ b/src/components/chat/HomeChatInput.tsx
@@ -0,0 +1,85 @@
+import { SendIcon, StopCircleIcon, X } from "lucide-react";
+import type React from "react";
+import { useEffect, useRef, useState } from "react";
+import { ModelPicker } from "@/components/ModelPicker";
+import { useSettings } from "@/hooks/useSettings";
+import { homeChatInputValueAtom } from "@/atoms/chatAtoms"; // Use a different atom for home input
+import { useAtom } from "jotai";
+import { useStreamChat } from "@/hooks/useStreamChat";
+
+export function HomeChatInput({ onSubmit }: { onSubmit: () => void }) {
+  const [inputValue, setInputValue] = useAtom(homeChatInputValueAtom);
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
+  const { settings, updateSettings, isAnyProviderSetup } = useSettings();
+  const { streamMessage, isStreaming, setIsStreaming } = useStreamChat(); // eslint-disable-line @typescript-eslint/no-unused-vars
+
+  const adjustHeight = () => {
+    const textarea = textareaRef.current;
+    if (textarea) {
+      textarea.style.height = "0px";
+      const scrollHeight = textarea.scrollHeight;
+      textarea.style.height = `${scrollHeight + 4}px`;
+    }
+  };
+
+  useEffect(() => {
+    adjustHeight();
+  }, [inputValue]);
+
+  const handleKeyPress = (e: React.KeyboardEvent) => {
+    if (e.key === "Enter" && !e.shiftKey) {
+      e.preventDefault();
+      onSubmit();
+    }
+  };
+
+  if (!settings) {
+    return null; // Or loading state
+  }
+
+  return (
+    <>
+      <div className="p-4">
+        <div className="flex flex-col space-y-2 border border-border rounded-lg bg-(--background-lighter) shadow-sm">
+          <div className="flex items-start space-x-2 ">
+            <textarea
+              ref={textareaRef}
+              value={inputValue}
+              onChange={(e) => setInputValue(e.target.value)}
+              onKeyPress={handleKeyPress}
+              placeholder="Ask Dyad to build..."
+              className="flex-1 p-2 focus:outline-none overflow-y-auto min-h-[40px] max-h-[200px]"
+              style={{ resize: "none" }}
+              disabled={isStreaming} // Should ideally reflect if *any* stream is happening
+            />
+            {isStreaming ? (
+              <button
+                className="px-2 py-2 mt-1 mr-2 text-(--sidebar-accent-fg) rounded-lg opacity-50 cursor-not-allowed" // Indicate disabled state
+                title="Cancel generation (unavailable here)"
+              >
+                <StopCircleIcon size={20} />
+              </button>
+            ) : (
+              <button
+                onClick={onSubmit}
+                disabled={!inputValue.trim() || !isAnyProviderSetup()}
+                className="px-2 py-2 mt-1 mr-2 hover:bg-(--background-darkest) text-(--sidebar-accent-fg) rounded-lg disabled:opacity-50"
+                title="Start new chat"
+              >
+                <SendIcon size={20} />
+              </button>
+            )}
+          </div>
+          <div className="px-2 pb-2">
+            <ModelPicker
+              selectedModel={settings.selectedModel}
+              onModelSelect={(model) =>
+                updateSettings({ selectedModel: model })
+              }
+            />
+          </div>
+        </div>
+      </div>
+    </>
+  );
+}
diff --git a/src/hooks/useProposal.ts b/src/hooks/useProposal.ts
index ec426f4..a10577b 100644
--- a/src/hooks/useProposal.ts
+++ b/src/hooks/useProposal.ts
@@ -1,30 +1,26 @@
 import { useState, useEffect, useCallback } from "react";
 import { IpcClient } from "@/ipc/ipc_client";
-import type { Proposal } from "@/lib/schemas"; // Import Proposal type
-
-// Define the structure returned by the IPC call
-interface ProposalResult {
-  proposal: Proposal;
-  messageId: number;
-}
-
+import type { Proposal, ProposalResult } from "@/lib/schemas"; // Import Proposal type
+import { proposalResultAtom } from "@/atoms/proposalAtoms";
+import { useAtom } from "jotai";
 export function useProposal(chatId?: number | undefined) {
-  const [proposalData, setProposalData] = useState<ProposalResult | null>(null);
+  const [proposalResult, setProposalResult] = useAtom(proposalResultAtom);
   const [isLoading, setIsLoading] = useState<boolean>(false);
   const [error, setError] = useState<string | null>(null);
 
   const fetchProposal = useCallback(
     async (innerChatId?: number) => {
       chatId = chatId ?? innerChatId;
+      console.log("fetching proposal for chatId", chatId);
       if (chatId === undefined) {
-        setProposalData(null);
+        setProposalResult(null);
         setIsLoading(false);
         setError(null);
         return;
       }
       setIsLoading(true);
       setError(null);
-      setProposalData(null); // Reset on new fetch
+      setProposalResult(null); // Reset on new fetch
       try {
         // Type assertion might be needed depending on how IpcClient is typed
         const result = (await IpcClient.getInstance().getProposal(
@@ -32,14 +28,14 @@ export function useProposal(chatId?: number | undefined) {
         )) as ProposalResult | null;
 
         if (result) {
-          setProposalData(result);
+          setProposalResult(result);
         } else {
-          setProposalData(null); // Explicitly set to null if IPC returns null
+          setProposalResult(null); // Explicitly set to null if IPC returns null
         }
       } catch (err: any) {
         console.error("Error fetching proposal:", err);
         setError(err.message || "Failed to fetch proposal");
-        setProposalData(null); // Clear proposal data on error
+        setProposalResult(null); // Clear proposal data on error
       } finally {
         setIsLoading(false);
       }
@@ -64,8 +60,8 @@ export function useProposal(chatId?: number | undefined) {
   );
 
   return {
-    proposal: proposalData?.proposal ?? null,
-    messageId: proposalData?.messageId,
+    proposal: proposalResult?.proposal ?? null,
+    messageId: proposalResult?.messageId,
     isLoading,
     error,
     refreshProposal, // Expose the refresh function
diff --git a/src/hooks/useStreamChat.ts b/src/hooks/useStreamChat.ts
index 41d1868..203de74 100644
--- a/src/hooks/useStreamChat.ts
+++ b/src/hooks/useStreamChat.ts
@@ -47,7 +47,6 @@ export function useStreamChat() {
       }
 
       setError(null);
-      console.log("streaming message - set messages", prompt);
       setMessages((currentMessages: Message[]) => {
         if (redo) {
           let remainingMessages = currentMessages.slice();
@@ -92,10 +91,10 @@ export function useStreamChat() {
             if (response.updatedFiles) {
               setIsPreviewOpen(true);
             }
+            refreshProposal(chatId);
 
             // Keep the same as below
             setIsStreaming(false);
-            refreshProposal(chatId);
             refreshChats();
             refreshApp();
             refreshVersions();
diff --git a/src/ipc/handlers/chat_stream_handlers.ts b/src/ipc/handlers/chat_stream_handlers.ts
index c53216a..adb983f 100644
--- a/src/ipc/handlers/chat_stream_handlers.ts
+++ b/src/ipc/handlers/chat_stream_handlers.ts
@@ -285,12 +285,12 @@ export function registerChatStreamHandlers() {
             chatId: req.chatId,
             updatedFiles: status.updatedFiles ?? false,
           } satisfies ChatResponseEnd);
+        } else {
+          event.sender.send("chat:response:end", {
+            chatId: req.chatId,
+            updatedFiles: false,
+          } satisfies ChatResponseEnd);
         }
-      } else {
-        event.sender.send("chat:response:end", {
-          chatId: req.chatId,
-          updatedFiles: false,
-        } satisfies ChatResponseEnd);
       }
 
       // Return the chat ID for backwards compatibility
diff --git a/src/ipc/ipc_client.ts b/src/ipc/ipc_client.ts
index a715faf..fcf55a4 100644
--- a/src/ipc/ipc_client.ts
+++ b/src/ipc/ipc_client.ts
@@ -17,15 +17,9 @@ import type {
   Message,
   Version,
 } from "./ipc_types";
-import type { Proposal } from "@/lib/schemas";
+import type { Proposal, ProposalResult } from "@/lib/schemas";
 import { showError } from "@/lib/toast";
 
-// Define the structure returned by getProposal
-interface ProposalResult {
-  proposal: Proposal;
-  messageId: number;
-}
-
 export interface ChatStreamCallbacks {
   onUpdate: (messages: Message[]) => void;
   onEnd: (response: ChatResponseEnd) => void;
diff --git a/src/lib/schemas.ts b/src/lib/schemas.ts
index fef934e..75cb620 100644
--- a/src/lib/schemas.ts
+++ b/src/lib/schemas.ts
@@ -117,3 +117,8 @@ export interface Proposal {
   securityRisks: SecurityRisk[];
   filesChanged: FileChange[];
 }
+
+export interface ProposalResult {
+  proposal: Proposal;
+  messageId: number;
+}
diff --git a/src/pages/home.tsx b/src/pages/home.tsx
index ad7f4cc..cad80b8 100644
--- a/src/pages/home.tsx
+++ b/src/pages/home.tsx
@@ -1,19 +1,18 @@
 import { useNavigate, useSearch } from "@tanstack/react-router";
 import { useAtom, useSetAtom } from "jotai";
-import { chatInputValueAtom } from "../atoms/chatAtoms";
+import { homeChatInputValueAtom } from "../atoms/chatAtoms";
 import { selectedAppIdAtom } from "@/atoms/appAtoms";
 import { IpcClient } from "@/ipc/ipc_client";
 import { generateCuteAppName } from "@/lib/utils";
 import { useLoadApps } from "@/hooks/useLoadApps";
 import { useSettings } from "@/hooks/useSettings";
 import { SetupBanner } from "@/components/SetupBanner";
-import { ChatInput } from "@/components/chat/ChatInput";
 import { isPreviewOpenAtom } from "@/atoms/viewAtoms";
 import { useState, useEffect } from "react";
 import { useStreamChat } from "@/hooks/useStreamChat";
-
+import { HomeChatInput } from "@/components/chat/HomeChatInput";
 export default function HomePage() {
-  const [inputValue, setInputValue] = useAtom(chatInputValueAtom);
+  const [inputValue, setInputValue] = useAtom(homeChatInputValueAtom);
   const navigate = useNavigate();
   const search = useSearch({ from: "/" });
   const setSelectedAppId = useSetAtom(selectedAppIdAtom);
@@ -87,7 +86,7 @@ export default function HomePage() {
       <SetupBanner />
 
       <div className="w-full">
-        <ChatInput onSubmit={handleSubmit} />
+        <HomeChatInput onSubmit={handleSubmit} />
 
         <div className="flex flex-wrap gap-4 mt-4">
           {[