diff --git a/pages/chat/@id/+Page.tsx b/pages/chat/@id/+Page.tsx
index d794680..cc0d951 100644
--- a/pages/chat/@id/+Page.tsx
+++ b/pages/chat/@id/+Page.tsx
@@ -28,6 +28,7 @@ import type {
   DraftMessage,
   OtherParameters,
   SendMessageStatus,
+  SendMessageStatusUI,
 } from "@/types";
 import Markdown from "react-markdown";
 import {
@@ -387,7 +388,7 @@ export default function ChatPage() {
         {
           signal: abortController.signal,
           onData: (data) => {
-            setSendMessageStatus(data);
+            setSendMessageStatus(data as SendMessageStatus);
 
             // If we've completed, update the UI and invalidate queries
             if (data.status === "completed") {
@@ -418,7 +419,7 @@ export default function ChatPage() {
             console.error("Subscription error:", error);
             setIsSendingMessage(false);
             setSendMessageStatus({
-              status: "error",
+              status: "error" as const,
               message: "An error occurred while sending the message",
             });
           },
@@ -434,7 +435,7 @@ export default function ChatPage() {
       console.error("Failed to start subscription:", error);
       setIsSendingMessage(false);
       setSendMessageStatus({
-        status: "error",
+        status: "error" as const,
         message: "Failed to start message sending process",
       });
     }
@@ -758,7 +759,9 @@ function Messages() {
         });
 
         /** Optimistically update the affected queries in react-query's cache: */
-        const previousMessages = await queryClient.getQueryData(
+        const previousMessages = await queryClient.getQueryData<
+          Array<CommittedMessage>
+        >(
           trpc.chat.messages.fetchByConversationId.queryKey({
             conversationId,
           })
@@ -772,7 +775,7 @@ function Messages() {
         const newMessages = previousMessages.filter(
           (m: CommittedMessage) => m.id !== messageIdToDelete
         );
-        queryClient.setQueryData(
+        queryClient.setQueryData<Array<CommittedMessage>>(
           trpc.chat.messages.fetchByConversationId.queryKey({
             conversationId,
           }),
@@ -878,7 +881,7 @@ function Messages() {
 }
 
 const StatusMessage = memo(
-  ({ sendMessageStatus }: { sendMessageStatus: SendMessageStatus | null }) => {
+  ({ sendMessageStatus }: { sendMessageStatus: SendMessageStatusUI }) => {
     const [displayMessage, setDisplayMessage] = useState(sendMessageStatus);
     const [isVisible, setIsVisible] = useState(sendMessageStatus !== null);
 
diff --git a/types.ts b/types.ts
index 0791fbc..c020552 100644
--- a/types.ts
+++ b/types.ts
@@ -1,19 +1,28 @@
 import type { UIMessage } from "ai";
 import type { generateText } from "ai";
 import type { Conversation, Fact, FactTrigger } from "./database/common";
+import type { chat } from "./server/trpc/chat";
 
 export type OtherParameters = Omit<
   Parameters<typeof generateText>[0],
-  "model" | "messages" | "abortSignal"
+  "model" | "messages" | "abortSignal" | "prompt"
 >;
 
 export type ConversationUI = Conversation & {};
 
-export type SendMessageStatus = {
-  status: string;
-  message: string;
-  result?: any;
-};
+/** Helper to extract the inner type of an async iterable. */
+type AsyncIterated<T> = T extends AsyncIterable<infer U> ? U : never;
+export type SendMessageStatus = AsyncIterated<
+  Awaited<ReturnType<typeof chat.sendMessage>>
+>;
+export type SendMessageStatusUI =
+  | SendMessageStatus
+  | {
+      readonly status: "error";
+      readonly message: string;
+      readonly result?: undefined;
+    }
+  | null;
 
 export type Store = {
   /** This is a string because Milvus sends it as a string, and the value
@@ -27,7 +36,7 @@ export type Store = {
   facts: Array<Fact>;
   factTriggers: Array<FactTrigger>;
   loading: boolean;
-  sendMessageStatus: SendMessageStatus | null;
+  sendMessageStatus: SendMessageStatusUI;
   isSendingMessage: boolean;
   setConversationId: (conversationId: string) => void;
   setConversationTitle: (conversationTitle: string) => void;
@@ -43,7 +52,7 @@ export type Store = {
   removeFact: (factId: string) => void;
   removeFactTrigger: (factTriggerId: string) => void;
   setLoading: (loading: boolean) => void;
-  setSendMessageStatus: (status: SendMessageStatus | null) => void;
+  setSendMessageStatus: (status: SendMessageStatusUI) => void;
   setIsSendingMessage: (isSending: boolean) => void;
 };