feat: perf optimizations

- improved search_knowledgebase_tool
- Added new endpoint to batch-fetch comments for multiple messages, reducing the number of API calls.
- Introduced CommentBatchRequest and CommentBatchResponse schemas for handling batch requests and responses.
- Updated chat_comments_service to validate message existence and permissions before fetching comments.
- Enhanced frontend with useBatchCommentsPreload hook to optimize comment loading for assistant messages.

surfsense_web/components/assistant-ui/thread.tsx

@@ -65,6 +65,7 @@ import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
 import { Button } from "@/components/ui/button";
 import { Spinner } from "@/components/ui/spinner";
 import type { Document } from "@/contracts/types/document.types";
+import { useBatchCommentsPreload } from "@/hooks/use-comments";
 import { useCommentsElectric } from "@/hooks/use-comments-electric";
 import { documentsApiService } from "@/lib/apis/documents-api.service";
 import { cn } from "@/lib/utils";
@@ -309,6 +310,22 @@ const Composer: FC = () => {
 	// Sync comments for the entire thread via Electric SQL (one subscription per thread)
 	useCommentsElectric(threadId);
 
+	// Batch-prefetch comments for all assistant messages so individual useComments
+	// hooks never fire their own network requests (eliminates N+1 API calls).
+	// Return a primitive string from the selector so useSyncExternalStore can
+	// compare snapshots by value and avoid infinite re-render loops.
+	const assistantIdsKey = useAssistantState(({ thread }) =>
+		thread.messages
+			.filter((m) => m.role === "assistant" && m.id?.startsWith("msg-"))
+			.map((m) => m.id!.replace("msg-", ""))
+			.join(",")
+	);
+	const assistantDbMessageIds = useMemo(
+		() => (assistantIdsKey ? assistantIdsKey.split(",").map(Number) : []),
+		[assistantIdsKey]
+	);
+	useBatchCommentsPreload(assistantDbMessageIds);
+
 	// Auto-focus editor on new chat page after mount
 	useEffect(() => {
 		if (isThreadEmpty && !hasAutoFocusedRef.current && editorRef.current) {

surfsense_web/contracts/types/chat-comments.types.ts

@@ -82,6 +82,22 @@ export const getCommentsResponse = z.object({
 	total_count: z.number(),
 });
 
+/**
+ * Batch-fetch comments for multiple messages
+ */
+export const getBatchCommentsRequest = z.object({
+	message_ids: z.array(z.number()).min(1).max(200),
+});
+
+export const commentListResponse = z.object({
+	comments: z.array(comment),
+	total_count: z.number(),
+});
+
+export const getBatchCommentsResponse = z.object({
+	comments_by_message: z.record(z.string(), commentListResponse),
+});
+
 /**
  * Create comment
  */
@@ -145,6 +161,8 @@ export type MentionComment = z.infer<typeof mentionComment>;
 export type Mention = z.infer<typeof mention>;
 export type GetCommentsRequest = z.infer<typeof getCommentsRequest>;
 export type GetCommentsResponse = z.infer<typeof getCommentsResponse>;
+export type GetBatchCommentsRequest = z.infer<typeof getBatchCommentsRequest>;
+export type GetBatchCommentsResponse = z.infer<typeof getBatchCommentsResponse>;
 export type CreateCommentRequest = z.infer<typeof createCommentRequest>;
 export type CreateCommentResponse = z.infer<typeof createCommentResponse>;
 export type CreateReplyRequest = z.infer<typeof createReplyRequest>;

surfsense_web/hooks/use-comments.ts

@@ -1,4 +1,5 @@
-import { useQuery } from "@tanstack/react-query";
+import { useQuery, useQueryClient } from "@tanstack/react-query";
+import { useEffect, useRef } from "react";
 import { chatCommentsApiService } from "@/lib/apis/chat-comments-api.service";
 import { cacheKeys } from "@/lib/query-client/cache-keys";
 
@@ -7,12 +8,84 @@ interface UseCommentsOptions {
 	enabled?: boolean;
 }
 
+// ---------------------------------------------------------------------------
+// Module-level coordination: when a batch request is in-flight, individual
+// useComments queryFns piggy-back on it instead of making their own requests.
+// ---------------------------------------------------------------------------
+let _batchInflight: Promise<void> | null = null;
+let _batchTargetIds = new Set<number>();
+
 export function useComments({ messageId, enabled = true }: UseCommentsOptions) {
+	const queryClient = useQueryClient();
+
 	return useQuery({
 		queryKey: cacheKeys.comments.byMessage(messageId),
 		queryFn: async () => {
+			// Yield one macro-task so the batch prefetch useEffect (which sets
+			// _batchInflight) has a chance to fire before we decide to fetch.
+			await new Promise<void>((r) => setTimeout(r, 0));
+
+			if (_batchInflight && _batchTargetIds.has(messageId)) {
+				await _batchInflight;
+				const cached = queryClient.getQueryData(cacheKeys.comments.byMessage(messageId));
+				if (cached) return cached;
+			}
+
 			return chatCommentsApiService.getComments({ message_id: messageId });
 		},
 		enabled: enabled && !!messageId,
+		staleTime: 30_000,
 	});
 }
+
+/**
+ * Batch-fetch comments for all given message IDs in a single request, then
+ * seed the per-message React Query cache so individual useComments hooks
+ * resolve from cache instead of firing their own requests.
+ */
+export function useBatchCommentsPreload(messageIds: number[]) {
+	const queryClient = useQueryClient();
+	const prevKeyRef = useRef<string>("");
+
+	useEffect(() => {
+		if (!messageIds.length) return;
+
+		const key = messageIds
+			.slice()
+			.sort((a, b) => a - b)
+			.join(",");
+		if (key === prevKeyRef.current) return;
+		prevKeyRef.current = key;
+
+		_batchTargetIds = new Set(messageIds);
+		let cancelled = false;
+
+		const promise = chatCommentsApiService
+			.getBatchComments({ message_ids: messageIds })
+			.then((data) => {
+				if (cancelled) return;
+				for (const [msgIdStr, commentList] of Object.entries(data.comments_by_message)) {
+					queryClient.setQueryData(cacheKeys.comments.byMessage(Number(msgIdStr)), commentList);
+				}
+			})
+			.catch(() => {
+				// Batch failed; individual queryFns will fall through to their own fetch
+			})
+			.finally(() => {
+				if (_batchInflight === promise) {
+					_batchInflight = null;
+					_batchTargetIds = new Set();
+				}
+			});
+
+		_batchInflight = promise;
+
+		return () => {
+			cancelled = true;
+			if (_batchInflight === promise) {
+				_batchInflight = null;
+				_batchTargetIds = new Set();
+			}
+		};
+	}, [messageIds, queryClient]);
+}

surfsense_web/lib/apis/chat-comments-api.service.ts

@@ -8,8 +8,11 @@ import {
 	type DeleteCommentRequest,
 	deleteCommentRequest,
 	deleteCommentResponse,
+	type GetBatchCommentsRequest,
 	type GetCommentsRequest,
 	type GetMentionsRequest,
+	getBatchCommentsRequest,
+	getBatchCommentsResponse,
 	getCommentsRequest,
 	getCommentsResponse,
 	getMentionsRequest,
@@ -22,6 +25,22 @@ import { ValidationError } from "@/lib/error";
 import { baseApiService } from "./base-api.service";
 
 class ChatCommentsApiService {
+	/**
+	 * Batch-fetch comments for multiple messages in one request
+	 */
+	getBatchComments = async (request: GetBatchCommentsRequest) => {
+		const parsed = getBatchCommentsRequest.safeParse(request);
+
+		if (!parsed.success) {
+			const errorMessage = parsed.error.issues.map((issue) => issue.message).join(", ");
+			throw new ValidationError(`Invalid request: ${errorMessage}`);
+		}
+
+		return baseApiService.post("/api/v1/messages/comments/batch", getBatchCommentsResponse, {
+			body: { message_ids: parsed.data.message_ids },
+		});
+	};
+
 	/**
 	 * Get comments for a message
 	 */