diff --git a/agent/uis/stock-price/index.tsx b/agent/uis/stock-price/index.tsx
index d773f07..2d7abf0 100644
--- a/agent/uis/stock-price/index.tsx
+++ b/agent/uis/stock-price/index.tsx
@@ -10,7 +10,7 @@ export default function StockPrice(props: {
   const [counter, setCounter] = useState(0);
 
   // useStream should be able to be infered from context
-  const thread = useStream<{ messages: Message[] }, { messages: Message[] }>({
+  const thread = useStream<{ messages: Message[] }>({
     assistantId: "assistant_123",
     apiUrl: "http://localhost:3123",
   });
@@ -22,7 +22,7 @@ export default function StockPrice(props: {
     .reverse()
     .find(
       (message): message is AIMessage =>
-        message.type === "ai" && !!message.tool_calls?.length,
+        message.type === "ai" && !!message.tool_calls?.length
     );
 
   const toolCallId = aiTool?.tool_calls?.[0]?.id;
diff --git a/src/providers/Stream.tsx b/src/providers/Stream.tsx
index c830637..0141c65 100644
--- a/src/providers/Stream.tsx
+++ b/src/providers/Stream.tsx
@@ -6,30 +6,24 @@ import type {
   RemoveUIMessage,
 } from "@langchain/langgraph-sdk/react-ui/types";
 
-type StreamContextType = ReturnType<
-  typeof useStream<
-    { messages: Message[]; ui: UIMessage[] },
-    {
+const useTypedStream = useStream<
+  { messages: Message[]; ui: UIMessage[] },
+  {
+    UpdateType: {
       messages?: Message[] | Message | string;
       ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
-    },
-    UIMessage | RemoveUIMessage
-  >
+    };
+    CustomType: UIMessage | RemoveUIMessage;
+  }
 >;
 
+type StreamContextType = ReturnType<typeof useTypedStream>;
 const StreamContext = createContext<StreamContextType | undefined>(undefined);
 
 export const StreamProvider: React.FC<{ children: ReactNode }> = ({
   children,
 }) => {
-  const streamValue = useStream<
-    { messages: Message[]; ui: UIMessage[] },
-    {
-      messages?: Message[] | Message | string;
-      ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
-    },
-    UIMessage | RemoveUIMessage
-  >({
+  const streamValue = useTypedStream({
     apiUrl: "http://localhost:2024",
     assistantId: "agent",
   });