feat: add Zod schemas for Surfsense docs and fix citation handling

surfsense_web/components/assistant-ui/markdown-text.tsx

@@ -20,7 +20,8 @@ const CITATION_REGEX = /\[citation:(doc-)?(\d+)\]/g;
 
 // Track chunk IDs to citation numbers mapping for consistent numbering
 // This map is reset when a new message starts rendering
-let chunkIdToCitationNumber: Map<number, number> = new Map();
+// Uses string keys to differentiate between doc and regular chunks (e.g., "doc-123" vs "123")
+let chunkIdToCitationNumber: Map<string, number> = new Map();
 let nextCitationNumber = 1;
 
 /**
@@ -37,11 +38,11 @@ export function resetCitationCounter() {
  */
 function getCitationNumber(chunkId: number, isDocsChunk: boolean): number {
 	const key = isDocsChunk ? `doc-${chunkId}` : String(chunkId);
-	const existingNumber = chunkIdToCitationNumber.get(key as unknown as number);
+	const existingNumber = chunkIdToCitationNumber.get(key);
 	if (existingNumber === undefined) {
-		chunkIdToCitationNumber.set(key as unknown as number, nextCitationNumber++);
+		chunkIdToCitationNumber.set(key, nextCitationNumber++);
 	}
-	return chunkIdToCitationNumber.get(key as unknown as number)!;
+	return chunkIdToCitationNumber.get(key)!;
 }
 
 /**

surfsense_web/components/new-chat/source-detail-panel.tsx

@@ -21,10 +21,16 @@ import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
 import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
 import { ScrollArea } from "@/components/ui/scroll-area";
+import type {
+	GetDocumentByChunkResponse,
+	GetSurfsenseDocsByChunkResponse,
+} from "@/contracts/types/document.types";
 import { documentsApiService } from "@/lib/apis/documents-api.service";
 import { cacheKeys } from "@/lib/query-client/cache-keys";
 import { cn } from "@/lib/utils";
 
+type DocumentData = GetDocumentByChunkResponse | GetSurfsenseDocsByChunkResponse;
+
 interface SourceDetailPanelProps {
 	open: boolean;
 	onOpenChange: (open: boolean) => void;
@@ -133,14 +139,16 @@ export function SourceDetailPanel({
 		data: documentData,
 		isLoading: isDocumentByChunkFetching,
 		error: documentByChunkFetchingError,
-	} = useQuery({
+	} = useQuery<DocumentData>({
 		queryKey: isDocsChunk
 			? cacheKeys.documents.byChunk(`doc-${chunkId}`)
 			: cacheKeys.documents.byChunk(chunkId.toString()),
-		queryFn: () =>
-			isDocsChunk
-				? documentsApiService.getSurfsenseDocByChunk(chunkId)
-				: documentsApiService.getDocumentByChunk({ chunk_id: chunkId }),
+		queryFn: async () => {
+			if (isDocsChunk) {
+				return documentsApiService.getSurfsenseDocByChunk(chunkId);
+			}
+			return documentsApiService.getDocumentByChunk({ chunk_id: chunkId });
+		},
 		enabled: !!chunkId && open,
 		staleTime: 5 * 60 * 1000,
 	});
@@ -332,7 +340,7 @@ export function SourceDetailPanel({
 									{documentData?.title || title || "Source Document"}
 								</h2>
 								<p className="text-sm text-muted-foreground mt-0.5">
-									{documentData
+									{documentData && "document_type" in documentData
 										? formatDocumentType(documentData.document_type)
 										: sourceType && formatDocumentType(sourceType)}
 									{documentData?.chunks && (
@@ -498,7 +506,8 @@ export function SourceDetailPanel({
 								<ScrollArea className="flex-1" ref={scrollAreaRef}>
 									<div className="p-6 lg:p-8 max-w-4xl mx-auto space-y-6">
 										{/* Document Metadata */}
-										{documentData.document_metadata &&
+										{"document_metadata" in documentData &&
+											documentData.document_metadata &&
 											Object.keys(documentData.document_metadata).length > 0 && (
 												<motion.div
 													initial={{ opacity: 0, y: 10 }}

surfsense_web/contracts/types/document.types.ts

@@ -59,6 +59,26 @@ export const documentWithChunks = document.extend({
 	),
 });
 
+/**
+ * Surfsense documentation schemas
+ * Follows the same pattern as document/documentWithChunks
+ */
+export const surfsenseDocsChunk = z.object({
+	id: z.number(),
+	content: z.string(),
+});
+
+export const surfsenseDocsDocument = z.object({
+	id: z.number(),
+	title: z.string(),
+	source: z.string(),
+	content: z.string(),
+});
+
+export const surfsenseDocsDocumentWithChunks = surfsenseDocsDocument.extend({
+	chunks: z.array(surfsenseDocsChunk),
+});
+
 /**
  * Get documents
  */
@@ -154,6 +174,15 @@ export const getDocumentByChunkRequest = z.object({
 
 export const getDocumentByChunkResponse = documentWithChunks;
 
+/**
+ * Get Surfsense docs by chunk
+ */
+export const getSurfsenseDocsByChunkRequest = z.object({
+	chunk_id: z.number(),
+});
+
+export const getSurfsenseDocsByChunkResponse = surfsenseDocsDocumentWithChunks;
+
 /**
  * Update document
  */
@@ -193,3 +222,8 @@ export type UpdateDocumentResponse = z.infer<typeof updateDocumentResponse>;
 export type DeleteDocumentRequest = z.infer<typeof deleteDocumentRequest>;
 export type DeleteDocumentResponse = z.infer<typeof deleteDocumentResponse>;
 export type DocumentTypeEnum = z.infer<typeof documentTypeEnum>;
+export type SurfsenseDocsChunk = z.infer<typeof surfsenseDocsChunk>;
+export type SurfsenseDocsDocument = z.infer<typeof surfsenseDocsDocument>;
+export type SurfsenseDocsDocumentWithChunks = z.infer<typeof surfsenseDocsDocumentWithChunks>;
+export type GetSurfsenseDocsByChunkRequest = z.infer<typeof getSurfsenseDocsByChunkRequest>;
+export type GetSurfsenseDocsByChunkResponse = z.infer<typeof getSurfsenseDocsByChunkResponse>;

surfsense_web/lib/apis/documents-api.service.ts

@@ -17,6 +17,7 @@ import {
 	getDocumentsResponse,
 	getDocumentTypeCountsRequest,
 	getDocumentTypeCountsResponse,
+	getSurfsenseDocsByChunkResponse,
 	type SearchDocumentsRequest,
 	searchDocumentsRequest,
 	searchDocumentsResponse,
@@ -214,10 +215,9 @@ class DocumentsApiService {
 	 * Used for resolving [citation:doc-XXX] citations
 	 */
 	getSurfsenseDocByChunk = async (chunkId: number) => {
-		// Response shape matches getDocumentByChunkResponse structure
 		return baseApiService.get(
 			`/api/v1/surfsense-docs/by-chunk/${chunkId}`,
-			getDocumentByChunkResponse
+			getSurfsenseDocsByChunkResponse
 		);
 	};