feat: integrate comments UI into chat messages

surfsense_backend/alembic/versions/66_add_chat_comments_table.py

@@ -26,15 +26,21 @@ def upgrade() -> None:
             content TEXT NOT NULL,
             created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
             updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
-        );
+        )
-
-        -- Create indexes
-        CREATE INDEX IF NOT EXISTS idx_chat_comments_message_id ON chat_comments(message_id);
-        CREATE INDEX IF NOT EXISTS idx_chat_comments_parent_id ON chat_comments(parent_id);
-        CREATE INDEX IF NOT EXISTS idx_chat_comments_author_id ON chat_comments(author_id);
-        CREATE INDEX IF NOT EXISTS idx_chat_comments_created_at ON chat_comments(created_at);
         """
     )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS idx_chat_comments_message_id ON chat_comments(message_id)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS idx_chat_comments_parent_id ON chat_comments(parent_id)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS idx_chat_comments_author_id ON chat_comments(author_id)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS idx_chat_comments_created_at ON chat_comments(created_at)"
+    )
 
 
 def downgrade() -> None:

surfsense_backend/alembic/versions/67_add_chat_comment_mentions_table.py

@@ -25,15 +25,15 @@ def upgrade() -> None:
             read BOOLEAN NOT NULL DEFAULT FALSE,
             created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
             UNIQUE (comment_id, mentioned_user_id)
-        );
+        )
-
-        -- Create indexes
-        CREATE INDEX IF NOT EXISTS idx_chat_comment_mentions_comment_id 
-            ON chat_comment_mentions(comment_id);
-        CREATE INDEX IF NOT EXISTS idx_chat_comment_mentions_user_unread 
-            ON chat_comment_mentions(mentioned_user_id) WHERE read = FALSE;
         """
     )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS idx_chat_comment_mentions_comment_id ON chat_comment_mentions(comment_id)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS idx_chat_comment_mentions_user_unread ON chat_comment_mentions(mentioned_user_id) WHERE read = FALSE"
+    )
 
 
 def downgrade() -> None:

surfsense_backend/app/schemas/chat_comments.py

@@ -55,6 +55,8 @@ class CommentReplyResponse(BaseModel):
     created_at: datetime
     updated_at: datetime
     is_edited: bool
+    can_edit: bool = False
+    can_delete: bool = False
 
     model_config = ConfigDict(from_attributes=True)
 
@@ -70,6 +72,8 @@ class CommentResponse(BaseModel):
     created_at: datetime
     updated_at: datetime
     is_edited: bool
+    can_edit: bool = False
+    can_delete: bool = False
     reply_count: int
     replies: list[CommentReplyResponse] = []
 

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

@@ -5,9 +5,11 @@ import {
 	MessagePrimitive,
 	useAssistantState,
 } from "@assistant-ui/react";
+import { useAtomValue } from "jotai";
 import { CheckIcon, CopyIcon, DownloadIcon, RefreshCwIcon } from "lucide-react";
 import type { FC } from "react";
-import { useContext } from "react";
+import { useContext, useState } from "react";
+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 {
@@ -16,6 +18,9 @@ import {
 } from "@/components/assistant-ui/thinking-steps";
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container";
+import { CommentTrigger } from "@/components/chat-comments/comment-trigger/comment-trigger";
+import { useComments } from "@/hooks/use-comments";
 
 export const MessageError: FC = () => {
 	return (
@@ -76,13 +81,53 @@ const AssistantMessageInner: FC = () => {
 	);
 };
 
+function parseMessageId(assistantUiMessageId: string | undefined): number | null {
+	if (!assistantUiMessageId) return null;
+	const match = assistantUiMessageId.match(/^msg-(\d+)$/);
+	return match ? Number.parseInt(match[1], 10) : null;
+}
+
 export const AssistantMessage: FC = () => {
+	const [isCommentPanelOpen, setIsCommentPanelOpen] = useState(false);
+	const messageId = useAssistantState(({ message }) => message?.id);
+	const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom);
+	const dbMessageId = parseMessageId(messageId);
+
+	const { data: commentsData } = useComments({
+		messageId: dbMessageId ?? 0,
+		enabled: !!dbMessageId,
+	});
+
+	const commentCount = commentsData?.total_count ?? 0;
+
 	return (
 		<MessagePrimitive.Root
-			className="aui-assistant-message-root fade-in slide-in-from-bottom-1 relative mx-auto w-full max-w-(--thread-max-width) animate-in py-3 duration-150"
+			className="aui-assistant-message-root group fade-in slide-in-from-bottom-1 relative mx-auto w-full max-w-(--thread-max-width) animate-in py-3 duration-150"
 			data-role="assistant"
 		>
 			<AssistantMessageInner />
+
+			{/* Comment trigger and floating panel */}
+			{dbMessageId && searchSpaceId && (
+				<div className="absolute -right-10 top-3">
+					<CommentTrigger
+						commentCount={commentCount}
+						isOpen={isCommentPanelOpen}
+						onClick={() => setIsCommentPanelOpen(!isCommentPanelOpen)}
+					/>
+					{isCommentPanelOpen && (
+						<div className="absolute left-full top-0 z-50 ml-2 animate-in fade-in slide-in-from-left-2 duration-200">
+							<CommentPanelContainer
+								messageId={dbMessageId}
+								searchSpaceId={Number(searchSpaceId)}
+								isOpen={true}
+								onClose={() => setIsCommentPanelOpen(false)}
+								maxHeight={400}
+							/>
+						</div>
+					)}
+				</div>
+			)}
 		</MessagePrimitive.Root>
 	);
 };

surfsense_web/components/assistant-ui/thread.tsx

@@ -36,6 +36,7 @@ import {
 	newLLMConfigsAtom,
 } from "@/atoms/new-llm-config/new-llm-config-query.atoms";
 import { currentUserAtom } from "@/atoms/user/user-query.atoms";
+import { AssistantMessage } from "@/components/assistant-ui/assistant-message";
 import { ComposerAddAttachment, ComposerAttachments } from "@/components/assistant-ui/attachment";
 import { ConnectorIndicator } from "@/components/assistant-ui/connector-popup";
 import {
@@ -590,17 +591,6 @@ const AssistantMessageInner: FC = () => {
 	);
 };
 
-const AssistantMessage: FC = () => {
-	return (
-		<MessagePrimitive.Root
-			className="aui-assistant-message-root fade-in slide-in-from-bottom-1 relative mx-auto w-full max-w-(--thread-max-width) animate-in py-3 duration-150"
-			data-role="assistant"
-		>
-			<AssistantMessageInner />
-		</MessagePrimitive.Root>
-	);
-};
-
 const AssistantActionBar: FC = () => {
 	return (
 		<ActionBarPrimitive.Root

surfsense_web/components/chat-comments/comment-composer/comment-composer.tsx

@@ -210,7 +210,7 @@ export function CommentComposer({
 
 	return (
 		<div className="flex flex-col gap-2">
-			<Popover open={mentionState.isActive} onOpenChange={(open) => !open && closeMentionPicker()}>
+			<Popover open={mentionState.isActive} onOpenChange={(open) => !open && closeMentionPicker()} modal={false}>
 				<PopoverAnchor asChild>
 					<Textarea
 						ref={textareaRef}

surfsense_web/components/chat-comments/comment-panel/comment-panel.tsx

@@ -45,7 +45,7 @@ export function CommentPanel({
 	const showEmptyState = !hasThreads && !isComposerOpen;
 
 	return (
-		<div className="flex w-96 flex-col rounded-lg border bg-card">
+		<div className="flex w-70 flex-col rounded-lg border bg-card">
 			{hasThreads && (
 				<div
 					className="overflow-y-auto"