Enhance chat session management in copilot-2 by implementing session creation and hydration logic. Refactor ChatContainer and EmptySession components to streamline user interactions and improve UI responsiveness. Update ChatInput to handle message sending with loading states, ensuring a smoother user experience.

2026-02-04 03:45:12 -05:00 · 2026-02-04 11:42:34 +05:30
parent df21b96fed
commit 6fce1f6084
8 changed files with 266 additions and 110 deletions
--- a/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/ChatContainer/ChatContainer.tsx
+++ b/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/ChatContainer/ChatContainer.tsx
@@ -3,79 +3,68 @@ import { UIDataTypes, UITools, UIMessage } from "ai";
 import { ChatMessagesContainer } from "../ChatMessagesContainer/ChatMessagesContainer";
 import { EmptySession } from "../EmptySession/EmptySession";
 import { ChatInput } from "@/components/contextual/Chat/components/ChatInput/ChatInput";
-import { postV2CreateSession } from "@/app/api/__generated__/endpoints/chat/chat";
-import { useState } from "react";
-import { parseAsString, useQueryState } from "nuqs";
 import { CopilotChatActionsProvider } from "../CopilotChatActionsProvider/CopilotChatActionsProvider";
+import { LayoutGroup, motion } from "framer-motion";

 export interface ChatContainerProps {
  messages: UIMessage<unknown, UIDataTypes, UITools>[];
  status: string;
  error: Error | undefined;
-  input: string;
-  setInput: (input: string) => void;
-  handleMessageSubmit: (e: React.FormEvent) => void;
-  onSend: (message: string) => void;
+  sessionId: string | null;
+  isCreatingSession: boolean;
+  onCreateSession: () => void | Promise<string>;
+  onSend: (message: string) => void | Promise<void>;
 }
 export const ChatContainer = ({
  messages,
  status,
  error,
-  input,
-  setInput,
-  handleMessageSubmit,
+  sessionId,
+  isCreatingSession,
+  onCreateSession,
  onSend,
 }: ChatContainerProps) => {
-  const [isCreating, setIsCreating] = useState(false);
-  const [sessionId, setSessionId] = useQueryState("sessionId", parseAsString);
-
-  async function createSession(e: React.FormEvent) {
-    e.preventDefault();
-    if (isCreating) return;
-    setIsCreating(true);
-    try {
-      const response = await postV2CreateSession({
-        body: JSON.stringify({}),
-      });
-      if (response.status === 200 && response.data?.id) {
-        setSessionId(response.data.id);
-      }
-    } finally {
-      setIsCreating(false);
-    }
-  }
+  const inputLayoutId = "copilot-2-chat-input";

  return (
    <CopilotChatActionsProvider onSend={onSend}>
-      <div className="mx-auto h-full w-full max-w-3xl pb-6">
-        <div className="flex h-full flex-col">
-          {sessionId ? (
-            <ChatMessagesContainer
-              messages={messages}
-              status={status}
-              error={error}
-              handleSubmit={handleMessageSubmit}
-              input={input}
-              setInput={setInput}
-            />
-          ) : (
-            <EmptySession
-              isCreating={isCreating}
-              onCreateSession={createSession}
-            />
-          )}
-          <div className="relative px-3 pt-2">
-            <div className="pointer-events-none absolute top-[-18px] z-10 h-6 w-full bg-gradient-to-b from-transparent to-[#f8f8f9]" />
-            <ChatInput
-              onSend={onSend}
-              disabled={status === "streaming" || !sessionId}
-              isStreaming={status === "streaming"}
-              onStop={() => {}}
-              placeholder="You can search or just ask"
-            />
+      <LayoutGroup id="copilot-2-chat-layout">
+        <div className="h-full w-full pb-6">
+          <div className="flex h-full flex-col">
+            {sessionId ? (
+              <div className="mx-auto flex h-full w-full max-w-3xl flex-col">
+                <ChatMessagesContainer
+                  messages={messages}
+                  status={status}
+                  error={error}
+                />
+                <motion.div
+                  layoutId={inputLayoutId}
+                  transition={{ type: "spring", bounce: 0.2, duration: 0.65 }}
+                  className="relative px-3 pt-2"
+                >
+                  <div className="pointer-events-none absolute top-[-18px] z-10 h-6 w-full bg-gradient-to-b from-transparent to-[#f8f8f9] dark:to-background" />
+                  <ChatInput
+                    inputId="chat-input-session"
+                    onSend={onSend}
+                    disabled={status === "streaming"}
+                    isStreaming={status === "streaming"}
+                    onStop={() => {}}
+                    placeholder="You can search or just ask"
+                  />
+                </motion.div>
+              </div>
+            ) : (
+              <EmptySession
+                inputLayoutId={inputLayoutId}
+                isCreatingSession={isCreatingSession}
+                onCreateSession={onCreateSession}
+                onSend={onSend}
+              />
+            )}
          </div>
        </div>
-      </div>
+      </LayoutGroup>
    </CopilotChatActionsProvider>
  );
 };
--- a/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/ChatMessagesContainer/ChatMessagesContainer.tsx
+++ b/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/ChatMessagesContainer/ChatMessagesContainer.tsx
@@ -24,9 +24,6 @@ interface ChatMessagesContainerProps {
  messages: UIMessage<unknown, UIDataTypes, UITools>[];
  status: string;
  error: Error | undefined;
-  handleSubmit: (e: React.FormEvent) => void;
-  input: string;
-  setInput: (input: string) => void;
 }

 export const ChatMessagesContainer = ({
--- a/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/CopilotChatActionsProvider/CopilotChatActionsProvider.tsx
+++ b/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/CopilotChatActionsProvider/CopilotChatActionsProvider.tsx
@@ -3,7 +3,7 @@
 import { CopilotChatActionsContext } from "./useCopilotChatActions";

 interface Props {
-  onSend: (message: string) => void;
+  onSend: (message: string) => void | Promise<void>;
  children: React.ReactNode;
 }

--- a/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/CopilotChatActionsProvider/useCopilotChatActions.ts
+++ b/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/CopilotChatActionsProvider/useCopilotChatActions.ts
@@ -3,7 +3,7 @@
 import { createContext, useContext } from "react";

 interface CopilotChatActions {
-  onSend: (message: string) => void;
+  onSend: (message: string) => void | Promise<void>;
 }

 const CopilotChatActionsContext = createContext<CopilotChatActions | null>(
--- a/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/EmptySession/EmptySession.tsx
+++ b/autogpt_platform/frontend/src/app/(platform)/copilot-2/components/EmptySession/EmptySession.tsx
@@ -1,23 +1,108 @@
+"use client";
+
+import {
+  getGreetingName,
+  getQuickActions,
+} from "@/app/(platform)/copilot/helpers";
+import { Button } from "@/components/atoms/Button/Button";
+import { Text } from "@/components/atoms/Text/Text";
+import { ChatInput } from "@/components/contextual/Chat/components/ChatInput/ChatInput";
+import { useSupabase } from "@/lib/supabase/hooks/useSupabase";
+import { SparkleIcon, SpinnerGapIcon } from "@phosphor-icons/react";
+import { motion } from "framer-motion";
+import { useState } from "react";
+
 interface Props {
-  isCreating: boolean;
-  onCreateSession: (e: React.FormEvent) => void;
+  inputLayoutId: string;
+  isCreatingSession: boolean;
+  onCreateSession: () => void | Promise<string>;
+  onSend: (message: string) => void | Promise<void>;
 }

-export function EmptySession({ isCreating, onCreateSession }: Props) {
+export function EmptySession({
+  inputLayoutId,
+  isCreatingSession,
+  onSend,
+}: Props) {
+  const { user } = useSupabase();
+  const greetingName = getGreetingName(user);
+  const quickActions = getQuickActions();
+  const [loadingAction, setLoadingAction] = useState<string | null>(null);
+
+  async function handleQuickActionClick(action: string) {
+    if (isCreatingSession || loadingAction) return;
+
+    setLoadingAction(action);
+    try {
+      await onSend(action);
+    } finally {
+      setLoadingAction(null);
+    }
+  }
+
  return (
-    <div className="flex h-full flex-1 flex-col items-center justify-center bg-zinc-100 p-4">
-      <h2 className="mb-4 text-xl font-semibold text-zinc-700">
-        Start a new conversation
-      </h2>
-      <form onSubmit={onCreateSession} className="w-full max-w-md">
-        <button
-          type="submit"
-          disabled={isCreating}
-          className="w-full rounded-md bg-blue-600 px-4 py-2 text-white transition-colors hover:bg-blue-700 disabled:opacity-50"
-        >
-          {isCreating ? "Creating..." : "Start New Chat"}
-        </button>
-      </form>
+    <div className="relative flex h-full flex-1 flex-col items-center justify-center overflow-hidden px-6 py-10 dark:bg-background">
+      <motion.div
+        className="relative w-full max-w-3xl"
+        initial={{ opacity: 0, y: 14, filter: "blur(6px)" }}
+        animate={{ opacity: 1, y: 0, filter: "blur(0px)" }}
+        transition={{ type: "spring", bounce: 0.2, duration: 0.7 }}
+      >
+        <div className="mx-auto flex flex-col items-center text-center">
+          <div className="mb-5 flex items-center gap-2 rounded-full border border-border/60 bg-white/70 px-3 py-1 text-xs text-muted-foreground shadow-sm backdrop-blur dark:bg-neutral-950/40">
+            <SparkleIcon className="h-3.5 w-3.5 text-purple-600" />
+            <span>Autopilot runs for you 24/7</span>
+          </div>
+
+          <Text variant="h3" className="mb-3 !text-[1.375rem] text-zinc-700">
+            Hey, <span className="text-violet-600">{greetingName}</span>
+          </Text>
+          <Text variant="h3" className="!font-normal">
+            What do you want to automate?
+          </Text>
+        </div>
+
+        <div className="mx-auto p-5 dark:bg-neutral-950/40">
+          <motion.div
+            layoutId={inputLayoutId}
+            transition={{ type: "spring", bounce: 0.2, duration: 0.65 }}
+            className="w-full"
+          >
+            <ChatInput
+              inputId="chat-input-empty"
+              onSend={onSend}
+              disabled={isCreatingSession}
+              placeholder='You can search or just ask - e.g. "create a blog post outline"'
+              className="w-full"
+            />
+          </motion.div>
+
+          <div className="mt-8 flex w-full flex-nowrap items-center justify-center gap-3 [-ms-overflow-style:none] [scrollbar-width:none] [&::-webkit-scrollbar]:hidden">
+            {quickActions.map((action) => (
+              <Button
+                key={action}
+                type="button"
+                variant="outline"
+                size="small"
+                onClick={() => void handleQuickActionClick(action)}
+                disabled={isCreatingSession || loadingAction !== null}
+                aria-busy={loadingAction === action}
+                leftIcon={
+                  loadingAction === action ? (
+                    <SpinnerGapIcon
+                      className="h-4 w-4 animate-spin"
+                      weight="bold"
+                    />
+                  ) : null
+                }
+                className="h-auto shrink-0 border-zinc-600 !px-4 !py-2 text-[1rem] text-zinc-600"
+              >
+                {action}
+              </Button>
+            ))}
+          </div>
+        </div>
+      </motion.div>
    </div>
  );
 }
--- a/autogpt_platform/frontend/src/app/(platform)/copilot-2/page.tsx
+++ b/autogpt_platform/frontend/src/app/(platform)/copilot-2/page.tsx
@@ -9,14 +9,24 @@ import { ChatContainer } from "./components/ChatContainer/ChatContainer";
 import { SidebarProvider, SidebarInset } from "@/components/ui/sidebar";
 import { Button } from "@/components/ui/button";
 import { CopyIcon, CheckIcon } from "@phosphor-icons/react";
-import { getV2GetSession } from "@/app/api/__generated__/endpoints/chat/chat";
+import {
+  getGetV2ListSessionsQueryKey,
+  getV2GetSession,
+  postV2CreateSession,
+} from "@/app/api/__generated__/endpoints/chat/chat";
 import { convertChatSessionMessagesToUiMessages } from "./helpers/convertChatSessionToUiMessages";
+import { useQueryClient } from "@tanstack/react-query";

 export default function Page() {
-  const [input, setInput] = useState("");
  const [copied, setCopied] = useState(false);
-  const [sessionId] = useQueryState("sessionId", parseAsString);
+  const [isCreatingSession, setIsCreatingSession] = useState(false);
+  const [sessionId, setSessionId] = useQueryState("sessionId", parseAsString);
  const hydrationSeq = useRef(0);
+  const lastHydratedSessionIdRef = useRef<string | null>(null);
+  const createSessionPromiseRef = useRef<Promise<string> | null>(null);
+  const queuedFirstMessageRef = useRef<string | null>(null);
+  const queuedFirstMessageResolverRef = useRef<(() => void) | null>(null);
+  const queryClient = useQueryClient();

  function handleCopySessionId() {
    if (!sessionId) return;
@@ -49,6 +59,41 @@ export default function Page() {
    transport: transport ?? undefined,
  });

+  const messagesRef = useRef(messages);
+
+  useEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
+
+  async function createSession(): Promise<string> {
+    if (sessionId) return sessionId;
+    if (createSessionPromiseRef.current) return createSessionPromiseRef.current;
+
+    setIsCreatingSession(true);
+    const promise = (async () => {
+      const response = await postV2CreateSession({
+        body: JSON.stringify({}),
+      });
+      if (response.status !== 200 || !response.data?.id) {
+        throw new Error("Failed to create chat session");
+      }
+      setSessionId(response.data.id);
+      queryClient.invalidateQueries({
+        queryKey: getGetV2ListSessionsQueryKey(),
+      });
+      return response.data.id;
+    })();
+
+    createSessionPromiseRef.current = promise;
+
+    try {
+      return await promise;
+    } finally {
+      createSessionPromiseRef.current = null;
+      setIsCreatingSession(false);
+    }
+  }
+
  useEffect(() => {
    hydrationSeq.current += 1;
    const seq = hydrationSeq.current;
@@ -56,11 +101,17 @@ export default function Page() {

    if (!sessionId) {
      setMessages([]);
+      lastHydratedSessionIdRef.current = null;
      return;
    }

    const currentSessionId = sessionId;

+    if (lastHydratedSessionIdRef.current !== currentSessionId) {
+      setMessages([]);
+      lastHydratedSessionIdRef.current = currentSessionId;
+    }
+
    async function hydrate() {
      try {
        const response = await getV2GetSession(currentSessionId, {
@@ -74,6 +125,13 @@ export default function Page() {
        );
        if (controller.signal.aborted) return;
        if (hydrationSeq.current !== seq) return;
+
+        const localMessagesCount = messagesRef.current.length;
+        const remoteMessagesCount = uiMessages.length;
+
+        if (remoteMessagesCount === 0) return;
+        if (localMessagesCount > remoteMessagesCount) return;
+
        setMessages(uiMessages);
      } catch (error) {
        if ((error as { name?: string } | null)?.name === "AbortError") return;
@@ -86,16 +144,33 @@ export default function Page() {
    return () => controller.abort();
  }, [sessionId, setMessages]);

-  function handleMessageSubmit(e: React.FormEvent) {
-    e.preventDefault();
-    if (!input.trim() || !sessionId) return;
+  useEffect(() => {
+    if (!sessionId) return;
+    const firstMessage = queuedFirstMessageRef.current;
+    if (!firstMessage) return;

-    sendMessage({ text: input });
-    setInput("");
-  }
+    queuedFirstMessageRef.current = null;
+    sendMessage({ text: firstMessage });
+    queuedFirstMessageResolverRef.current?.();
+    queuedFirstMessageResolverRef.current = null;
+  }, [sendMessage, sessionId]);

-  function onSend(message: string) {
-    sendMessage({ text: message });
+  async function onSend(message: string) {
+    const trimmed = message.trim();
+    if (!trimmed) return;
+
+    if (sessionId) {
+      sendMessage({ text: trimmed });
+      return;
+    }
+
+    queuedFirstMessageRef.current = trimmed;
+    const sentPromise = new Promise<void>((resolve) => {
+      queuedFirstMessageResolverRef.current = resolve;
+    });
+
+    await createSession();
+    await sentPromise;
  }

  return (
@@ -104,7 +179,7 @@ export default function Page() {
      className="h-[calc(100vh-72px)] min-h-0"
    >
      <ChatSidebar />
-      <SidebarInset className="relative flex h-[calc(100vh-80px)] flex-col">
+      <SidebarInset className="relative flex h-[calc(100vh-80px)] flex-col overflow-hidden ring-1 ring-zinc-300">
        {sessionId && (
          <div className="absolute flex items-center px-4 py-4">
            <div className="flex items-center gap-2 rounded-3xl border border-neutral-400 bg-neutral-100 px-3 py-1.5 text-sm text-neutral-600 dark:bg-neutral-800 dark:text-neutral-400">
@@ -129,9 +204,9 @@ export default function Page() {
            messages={messages}
            status={status}
            error={error}
-            input={input}
-            setInput={setInput}
-            handleMessageSubmit={handleMessageSubmit}
+            sessionId={sessionId}
+            isCreatingSession={isCreatingSession}
+            onCreateSession={createSession}
            onSend={onSend}
          />
        </div>
--- a/autogpt_platform/frontend/src/components/contextual/Chat/components/ChatInput/ChatInput.tsx
+++ b/autogpt_platform/frontend/src/components/contextual/Chat/components/ChatInput/ChatInput.tsx
@@ -11,12 +11,13 @@ import { useChatInput } from "./useChatInput";
 import { useVoiceRecording } from "./useVoiceRecording";

 export interface Props {
-  onSend: (message: string) => void;
+  onSend: (message: string) => void | Promise<void>;
  disabled?: boolean;
  isStreaming?: boolean;
  onStop?: () => void;
  placeholder?: string;
  className?: string;
+  inputId?: string;
 }

 export function ChatInput({
@@ -26,6 +27,7 @@ export function ChatInput({
  onStop,
  placeholder = "Type your message...",
  className,
+  inputId = "chat-input",
 }: Props) {
  const inputId = "chat-input";
  const {
--- a/autogpt_platform/frontend/src/components/contextual/Chat/components/ChatInput/useChatInput.ts
+++ b/autogpt_platform/frontend/src/components/contextual/Chat/components/ChatInput/useChatInput.ts
@@ -21,6 +21,7 @@ export function useChatInput({
 }: Args) {
  const [value, setValue] = useState("");
  const [hasMultipleLines, setHasMultipleLines] = useState(false);
+  const [isSending, setIsSending] = useState(false);

  useEffect(
    function focusOnMount() {
@@ -100,34 +101,40 @@ export function useChatInput({
    }
  }, [value, maxRows, inputId]);

-  const handleSend = () => {
-    if (disabled || !value.trim()) return;
-    onSend(value.trim());
-    setValue("");
-    setHasMultipleLines(false);
-    const textarea = document.getElementById(inputId) as HTMLTextAreaElement;
-    const wrapper = document.getElementById(
-      `${inputId}-wrapper`,
-    ) as HTMLDivElement;
-    if (textarea) {
-      textarea.style.height = "auto";
+  async function handleSend() {
+    if (disabled || isSending || !value.trim()) return;
+
+    setIsSending(true);
+    try {
+      await onSend(value.trim());
+      setValue("");
+      setHasMultipleLines(false);
+      const textarea = document.getElementById(inputId) as HTMLTextAreaElement;
+      const wrapper = document.getElementById(
+        `${inputId}-wrapper`,
+      ) as HTMLDivElement;
+      if (textarea) {
+        textarea.style.height = "auto";
+      }
+      if (wrapper) {
+        wrapper.style.height = "";
+        wrapper.style.maxHeight = "";
+      }
+    } finally {
+      setIsSending(false);
    }
-    if (wrapper) {
-      wrapper.style.height = "";
-      wrapper.style.maxHeight = "";
-    }
-  };
+  }

  function handleKeyDown(event: KeyboardEvent<HTMLTextAreaElement>) {
    if (event.key === "Enter" && !event.shiftKey) {
      event.preventDefault();
-      handleSend();
+      void handleSend();
    }
  }

  function handleSubmit(e: FormEvent<HTMLFormElement>) {
    e.preventDefault();
-    handleSend();
+    void handleSend();
  }

  function handleChange(e: ChangeEvent<HTMLTextAreaElement>) {
@@ -142,5 +149,6 @@ export function useChatInput({
    handleSubmit,
    handleChange,
    hasMultipleLines,
+    isSending,
  };
 }