fix: Only create thrad id at request time (#419)

Author: bracesproul

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+  "python.languageServer": "None"
+}

frontend/app/components/ChatLangChain.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import React, { useState } from "react";
+import React, { useEffect, useRef, useState } from "react";
 import {
   AppendMessage,
   AssistantRuntimeProvider,
@@ -19,37 +19,74 @@ import { SelectModel } from "./SelectModel";
 import { ThreadHistory } from "./thread-history";
 import { Toaster } from "./ui/toaster";
 import { useGraphContext } from "../contexts/GraphContext";
+import { useQueryState } from "nuqs";
 
 function ChatLangChainComponent(): React.ReactElement {
   const { toast } = useToast();
   const { threadsData, userData, graphData } = useGraphContext();
   const { userId } = userData;
-  const { getUserThreads, threadId: currentThread } = threadsData;
-  const { messages, setMessages, streamMessage } = graphData;
+  const { getUserThreads, createThread, getThreadById } = threadsData;
+  const { messages, setMessages, streamMessage, switchSelectedThread } =
+    graphData;
   const [isRunning, setIsRunning] = useState(false);
+  const [threadId, setThreadId] = useQueryState("threadId");
 
-  const isSubmitDisabled = !userId || !currentThread;
+  const hasCheckedThreadIdParam = useRef(false);
+  useEffect(() => {
+    if (typeof window === "undefined" || hasCheckedThreadIdParam.current)
+      return;
+    if (!threadId) {
+      hasCheckedThreadIdParam.current = true;
+      return;
+    }
+
+    hasCheckedThreadIdParam.current = true;
+
+    try {
+      getThreadById(threadId).then((thread) => {
+        if (!thread) {
+          setThreadId(null);
+          return;
+        }
+
+        switchSelectedThread(thread);
+      });
+    } catch (e) {
+      console.error("Failed to fetch thread in query param", e);
+      setThreadId(null);
+    }
+  }, [threadId]);
+
+  const isSubmitDisabled = !userId;
 
   async function onNew(message: AppendMessage): Promise<void> {
     if (isSubmitDisabled) {
-      let description = "";
-      if (!userId) {
-        description = "Unable to find user ID. Please try again later.";
-      } else if (!currentThread) {
-        description =
-          "Unable to find current thread ID. Please try again later.";
-      }
       toast({
         title: "Failed to send message",
-        description,
+        description: "Unable to find user ID. Please try again later.",
       });
       return;
     }
     if (message.content[0]?.type !== "text") {
       throw new Error("Only text messages are supported");
     }
+
     setIsRunning(true);
 
+    let currentThreadId = threadId;
+    if (!currentThreadId) {
+      const thread = await createThread(userId);
+      if (!thread) {
+        toast({
+          title: "Error",
+          description: "Thread creation failed.",
+        });
+        return;
+      }
+      setThreadId(thread.thread_id);
+      currentThreadId = thread.thread_id;
+    }
+
     try {
       const humanMessage = new HumanMessage({
         content: message.content[0].text,
@@ -58,7 +95,7 @@ function ChatLangChainComponent(): React.ReactElement {
 
       setMessages((prevMessages) => [...prevMessages, humanMessage]);
 
-      await streamMessage({
+      await streamMessage(currentThreadId, {
         messages: [convertToOpenAIFormat(humanMessage)],
       });
     } finally {

frontend/app/components/thread-history/index.tsx

@@ -1,29 +1,21 @@
 import { TooltipIconButton } from "../ui/assistant-ui/tooltip-icon-button";
 import { SquarePen, History } from "lucide-react";
 import { Sheet, SheetContent, SheetTrigger } from "../ui/sheet";
-import { useToast } from "../../hooks/use-toast";
 import { Skeleton } from "../ui/skeleton";
 import React from "react";
 import { useGraphContext } from "../../contexts/GraphContext";
 import { groupThreads } from "./utils";
 import { ThreadsList } from "./thread-list";
+import { useQueryState } from "nuqs";
 
 const LoadingThread = () => <Skeleton className="w-full h-8 bg-[#373737]" />;
 
 function ThreadHistoryComponent() {
-  const { toast } = useToast();
   const { threadsData, userData, graphData } = useGraphContext();
-  const {
-    userThreads,
-    isUserThreadsLoading,
-    threadId,
-    deleteThread,
-    createThread,
-    getUserThreads,
-  } = threadsData;
+  const { userThreads, isUserThreadsLoading, deleteThread } = threadsData;
   const { userId } = userData;
-  const { messages, switchSelectedThread, setMessages } = graphData;
-  const isEmpty = messages.length === 0;
+  const { switchSelectedThread, setMessages } = graphData;
+  const [_threadId, setThreadId] = useQueryState("threadId");
 
   const clearMessages = () => {
     setMessages([]);
@@ -40,25 +32,9 @@ function ThreadHistoryComponent() {
     deleteThreadAndClearMessages,
   );
 
-  const createAndSetupNewThread = async () => {
-    if (!userId) {
-      toast({
-        title: "Error creating thread",
-        description: "Your user ID was not found. Please try again later.",
-      });
-      return;
-    }
-    const currentThread = userThreads.find(
-      (thread) => thread.thread_id === threadId,
-    );
-    if (currentThread && !currentThread.values && isEmpty) {
-      return;
-    }
-
+  const createNewSession = async () => {
+    setThreadId(null);
     clearMessages();
-    await createThread(userId);
-    // Re-fetch threads so that the new thread shows up.
-    await getUserThreads(userId);
   };
 
   return (
@@ -73,7 +49,7 @@ function ThreadHistoryComponent() {
                 tooltip="New chat"
                 variant="ghost"
                 className="w-fit p-2"
-                onClick={createAndSetupNewThread}
+                onClick={createNewSession}
               >
                 <SquarePen className="w-5 h-5" />
               </TooltipIconButton>
@@ -121,7 +97,7 @@ function ThreadHistoryComponent() {
             tooltip="New chat"
             variant="ghost"
             className="w-fit h-fit p-2"
-            onClick={createAndSetupNewThread}
+            onClick={createNewSession}
           >
             <SquarePen className="w-6 h-6" />
           </TooltipIconButton>

frontend/app/contexts/GraphContext.tsx

@@ -1,11 +1,12 @@
+"use client";
+
 import { parsePartialJson } from "@langchain/core/output_parsers";
 import {
   createContext,
   Dispatch,
   ReactNode,
   SetStateAction,
   useContext,
-  useEffect,
   useState,
 } from "react";
 import { AIMessage, BaseMessage, HumanMessage } from "@langchain/core/messages";
@@ -18,14 +19,15 @@ import { useRuns } from "../hooks/useRuns";
 import { useUser } from "../hooks/useUser";
 import { addDocumentLinks, createClient, nodeToStep } from "./utils";
 import { Thread } from "@langchain/langgraph-sdk";
+import { useQueryState } from "nuqs";
 
 interface GraphData {
   messages: BaseMessage[];
   selectedModel: ModelOptions;
   setSelectedModel: Dispatch<SetStateAction<ModelOptions>>;
   setMessages: Dispatch<SetStateAction<BaseMessage[]>>;
-  streamMessage: (params: GraphInput) => Promise<void>;
-  switchSelectedThread: (thread: Thread) => Promise<void>;
+  streamMessage: (currentThreadId: string, params: GraphInput) => Promise<void>;
+  switchSelectedThread: (thread: Thread) => void;
 }
 
 type UserDataContextType = ReturnType<typeof useUser>;
@@ -49,30 +51,31 @@ export function GraphProvider({ children }: { children: ReactNode }) {
   const {
     isUserThreadsLoading,
     userThreads,
-    setUserThreads,
     getThreadById,
+    setUserThreads,
     getUserThreads,
     createThread,
-    searchOrCreateThread,
-    threadId,
-    setThreadId,
     deleteThread,
   } = useThreads(userId);
   const { toast } = useToast();
   const { shareRun } = useRuns();
   const [messages, setMessages] = useState<BaseMessage[]>([]);
-  const [selectedModel, setSelectedModel] =
-    useState<ModelOptions>("anthropic/claude-3-5-haiku-20241022");
+  const [selectedModel, setSelectedModel] = useState<ModelOptions>(
+    "anthropic/claude-3-5-haiku-20241022",
+  );
+  const [_threadId, setThreadId] = useQueryState("threadId");
 
-  const streamMessage = async (params: GraphInput): Promise<void> => {
-    if (!threadId) {
+  const streamMessage = async (
+    currentThreadId: string,
+    params: GraphInput,
+  ): Promise<void> => {
+    if (!userId) {
       toast({
         title: "Error",
-        description: "Thread ID not found",
+        description: "User ID not found",
       });
       return;
     }
-
     const client = createClient();
 
     const input = {
@@ -92,7 +95,7 @@ export function GraphProvider({ children }: { children: ReactNode }) {
       }),
     };
 
-    const stream = client.runs.stream(threadId, "chat", {
+    const stream = client.runs.stream(currentThreadId, "chat", {
       input,
       streamMode: "events",
       config: {
@@ -556,7 +559,7 @@ export function GraphProvider({ children }: { children: ReactNode }) {
     }
   };
 
-  const switchSelectedThread = async (thread: Thread) => {
+  const switchSelectedThread = (thread: Thread) => {
     setThreadId(thread.thread_id);
     if (!thread.values) {
       setMessages([]);
@@ -658,13 +661,10 @@ export function GraphProvider({ children }: { children: ReactNode }) {
     threadsData: {
       isUserThreadsLoading,
       userThreads,
-      setUserThreads,
       getThreadById,
+      setUserThreads,
       getUserThreads,
       createThread,
-      searchOrCreateThread,
-      threadId,
-      setThreadId,
       deleteThread,
     },
     graphData: {

frontend/app/hooks/useThreads.tsx

@@ -1,9 +1,10 @@
+"use client";
+
 import { useEffect, useState } from "react";
 
 import { Client, Thread } from "@langchain/langgraph-sdk";
-import { getCookie, setCookie } from "../utils/cookies";
 import { useToast } from "./use-toast";
-import { THREAD_ID_COOKIE_NAME } from "../utils/constants";
+import { useQueryState } from "nuqs";
 
 export const createClient = () => {
   const apiUrl = process.env.NEXT_PUBLIC_API_URL ?? "http://localhost:3000/api";
@@ -16,38 +17,13 @@ export function useThreads(userId: string | undefined) {
   const { toast } = useToast();
   const [isUserThreadsLoading, setIsUserThreadsLoading] = useState(false);
   const [userThreads, setUserThreads] = useState<Thread[]>([]);
-  const [threadId, setThreadId] = useState<string>();
+  const [threadId, setThreadId] = useQueryState("threadId");
 
   useEffect(() => {
     if (typeof window == "undefined" || !userId) return;
     getUserThreads(userId);
   }, [userId]);
 
-  useEffect(() => {
-    if (threadId || typeof window === "undefined" || !userId) return;
-    searchOrCreateThread(userId);
-  }, [userId]);
-
-  const searchOrCreateThread = async (id: string) => {
-    const threadIdCookie = getCookie(THREAD_ID_COOKIE_NAME);
-    if (!threadIdCookie) {
-      await createThread(id);
-      return;
-    }
-    // Thread ID is in cookies.
-
-    const thread = await getThreadById(threadIdCookie);
-    if (!thread.values) {
-      // No values = no activity. Can keep.
-      setThreadId(threadIdCookie);
-      return;
-    } else {
-      // Current thread has activity. Create a new thread.
-      await createThread(id);
-      return;
-    }
-  };
-
   const createThread = async (id: string) => {
     const client = createClient();
     let thread;
@@ -61,7 +37,6 @@ export function useThreads(userId: string | undefined) {
         throw new Error("Thread creation failed.");
       }
       setThreadId(thread.thread_id);
-      setCookie(THREAD_ID_COOKIE_NAME, thread.thread_id);
     } catch (e) {
       console.error("Error creating thread", e);
       toast({
@@ -111,28 +86,24 @@ export function useThreads(userId: string | undefined) {
       );
       return newThreads;
     });
+    const client = createClient();
+    await client.threads.delete(id);
     if (id === threadId) {
+      // Remove the threadID from query params, and refetch threads to
+      // update the sidebar UI.
       clearMessages();
-      // Create a new thread. Use .then to avoid blocking the UI.
-      // Once completed re-fetch threads to update UI.
-      createThread(userId).then(async () => {
-        await getUserThreads(userId);
-      });
+      getUserThreads(userId);
+      setThreadId(null);
     }
-    const client = createClient();
-    await client.threads.delete(id);
   };
 
   return {
     isUserThreadsLoading,
     userThreads,
-    setUserThreads,
     getThreadById,
+    setUserThreads,
     getUserThreads,
     createThread,
-    searchOrCreateThread,
-    threadId,
-    setThreadId,
     deleteThread,
   };
 }