Remove branch picker and restrict edit/reload to last message

surfsense_web/components/assistant-ui/assistant-message.tsx

@@ -4,20 +4,19 @@ import {
 	ErrorPrimitive,
 	MessagePrimitive,
 	useAssistantState,
+	useMessage,
 } from "@assistant-ui/react";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
+import { useAtom, useAtomValue } from "jotai";
 import { CheckIcon, CopyIcon, DownloadIcon, MessageSquare, RefreshCwIcon } from "lucide-react";
 import type { FC } from "react";
 import { useContext, useEffect, useMemo, useRef, useState } from "react";
 import {
 	addingCommentToMessageIdAtom,
-	clearTargetCommentIdAtom,
 	commentsCollapsedAtom,
 	commentsEnabledAtom,
 	targetCommentIdAtom,
 } from "@/atoms/chat/current-thread.atom";
 import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
-import { BranchPicker } from "@/components/assistant-ui/branch-picker";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import {
 	ThinkingStepsContext,
@@ -84,7 +83,6 @@ const AssistantMessageInner: FC = () => {
 			</div>
 
 			<div className="aui-assistant-message-footer mt-1 mb-5 ml-2 flex">
-				<BranchPicker />
 				<AssistantActionBar />
 			</div>
 		</>
@@ -126,7 +124,6 @@ export const AssistantMessage: FC = () => {
 
 	// Target comment navigation - read target from global atom
 	const targetCommentId = useAtomValue(targetCommentIdAtom);
-	const clearTargetCommentId = useSetAtom(clearTargetCommentIdAtom);
 
 	// Check if target comment belongs to this message (including replies)
 	const hasTargetComment = useMemo(() => {
@@ -263,6 +260,8 @@ export const AssistantMessage: FC = () => {
 };
 
 const AssistantActionBar: FC = () => {
+	const { isLast } = useMessage();
+
 	return (
 		<ActionBarPrimitive.Root
 			hideWhenRunning
@@ -285,11 +284,14 @@ const AssistantActionBar: FC = () => {
 					<DownloadIcon />
 				</TooltipIconButton>
 			</ActionBarPrimitive.ExportMarkdown>
-			<ActionBarPrimitive.Reload asChild>
+			{/* Only allow regenerating the last assistant message */}
-				<TooltipIconButton tooltip="Refresh">
+			{isLast && (
-					<RefreshCwIcon />
+				<ActionBarPrimitive.Reload asChild>
-				</TooltipIconButton>
+					<TooltipIconButton tooltip="Refresh">
-			</ActionBarPrimitive.Reload>
+						<RefreshCwIcon />
+					</TooltipIconButton>
+				</ActionBarPrimitive.Reload>
+			)}
 		</ActionBarPrimitive.Root>
 	);
 };

surfsense_web/components/assistant-ui/branch-picker.tsx

@@ -1,32 +0,0 @@
-import { BranchPickerPrimitive } from "@assistant-ui/react";
-import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
-import type { FC } from "react";
-import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
-import { cn } from "@/lib/utils";
-
-export const BranchPicker: FC<BranchPickerPrimitive.Root.Props> = ({ className, ...rest }) => {
-	return (
-		<BranchPickerPrimitive.Root
-			hideWhenSingleBranch
-			className={cn(
-				"aui-branch-picker-root -ml-2 mr-2 inline-flex items-center text-muted-foreground text-xs",
-				className
-			)}
-			{...rest}
-		>
-			<BranchPickerPrimitive.Previous asChild>
-				<TooltipIconButton tooltip="Previous">
-					<ChevronLeftIcon />
-				</TooltipIconButton>
-			</BranchPickerPrimitive.Previous>
-			<span className="aui-branch-picker-state font-medium">
-				<BranchPickerPrimitive.Number /> / <BranchPickerPrimitive.Count />
-			</span>
-			<BranchPickerPrimitive.Next asChild>
-				<TooltipIconButton tooltip="Next">
-					<ChevronRightIcon />
-				</TooltipIconButton>
-			</BranchPickerPrimitive.Next>
-		</BranchPickerPrimitive.Root>
-	);
-};

surfsense_web/components/assistant-ui/user-message.tsx

@@ -4,7 +4,6 @@ import { FileText, PencilIcon } from "lucide-react";
 import { type FC, useState } from "react";
 import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom";
 import { UserMessageAttachments } from "@/components/assistant-ui/attachment";
-import { BranchPicker } from "@/components/assistant-ui/branch-picker";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 
 interface AuthorMetadata {
@@ -95,24 +94,47 @@ export const UserMessage: FC = () => {
 					</div>
 				)}
 			</div>
-
-			<BranchPicker className="aui-user-branch-picker -mr-1 col-span-full col-start-1 row-start-3 justify-end" />
 		</MessagePrimitive.Root>
 	);
 };
 
 const UserActionBar: FC = () => {
+	const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
+	
+	// Get current message ID
+	const currentMessageId = useAssistantState(({ message }) => message?.id);
+	
+	// Find the last user message ID in the thread (computed once, memoized by selector)
+	const lastUserMessageId = useAssistantState(({ thread }) => {
+		const messages = thread.messages;
+		for (let i = messages.length - 1; i >= 0; i--) {
+			if (messages[i].role === "user") {
+				return messages[i].id;
+			}
+		}
+		return null;
+	});
+
+	// Simple comparison - no iteration needed per message
+	const isLastUserMessage = currentMessageId === lastUserMessageId;
+	
+	// Show edit button only on the last user message and when thread is not running
+	const canEdit = isLastUserMessage && !isThreadRunning;
+
 	return (
 		<ActionBarPrimitive.Root
 			hideWhenRunning
 			autohide="not-last"
 			className="aui-user-action-bar-root flex flex-col items-end"
 		>
-			<ActionBarPrimitive.Edit asChild>
+			{/* Only allow editing the last user message */}
-				<TooltipIconButton tooltip="Edit" className="aui-user-action-edit p-4">
+			{canEdit && (
-					<PencilIcon />
+				<ActionBarPrimitive.Edit asChild>
-				</TooltipIconButton>
+					<TooltipIconButton tooltip="Edit" className="aui-user-action-edit p-4">
-			</ActionBarPrimitive.Edit>
+						<PencilIcon />
+					</TooltipIconButton>
+				</ActionBarPrimitive.Edit>
+			)}
 		</ActionBarPrimitive.Root>
 	);
 };