From 2d134439ecae7a2f1492e7880429617edff5f222 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 26 May 2026 21:52:04 +0530 Subject: [PATCH] feat(web): enhance mention handling to support connectors and improve document key management --- surfsense_backend/app/schemas/new_chat.py | 19 ++--- .../app/tasks/chat/persistence.py | 19 +++-- .../new-chat/[[...chat_id]]/page.tsx | 15 ++-- .../atoms/chat/mentioned-documents.atom.ts | 34 +++----- .../assistant-ui/inline-mention-editor.tsx | 68 +++++++++++----- .../components/assistant-ui/thread.tsx | 31 ++++---- .../components/assistant-ui/user-message.tsx | 2 +- .../components/documents/FolderTreeView.tsx | 1 - .../new-chat/document-mention-picker.tsx | 44 ++++++----- surfsense_web/lib/chat/mention-doc-key.ts | 12 +-- surfsense_web/lib/connector-telemetry.ts | 79 +++++++------------ 11 files changed, 160 insertions(+), 164 deletions(-) diff --git a/surfsense_backend/app/schemas/new_chat.py b/surfsense_backend/app/schemas/new_chat.py index c721f495e..8b49413c6 100644 --- a/surfsense_backend/app/schemas/new_chat.py +++ b/surfsense_backend/app/schemas/new_chat.py @@ -203,13 +203,11 @@ class NewChatUserImagePart(BaseModel): class MentionedDocumentInfo(BaseModel): """Display metadata for a single ``@``-mention chip. - Carries either a knowledge-base document or a knowledge-base folder - (discriminated by ``kind``). The full triple - ``{id, title, document_type}`` is forwarded by the frontend mention - chip so the server can embed it in the persisted user message - ``ContentPart[]`` (single ``mentioned-documents`` part). The - history loader then renders the chips on reload without an extra - fetch — mirrors the pre-refactor frontend ``persistUserTurn`` shape. + Carries a knowledge-base document, knowledge-base folder, or + connected account (discriminated by ``kind``). Each kind uses its + real identity fields: docs carry ``document_type``, folders carry + only their folder id/title, and connectors carry ``connector_type`` + plus account metadata. ``kind`` defaults to ``"doc"`` so legacy clients and persisted rows that predate folder mentions deserialise unchanged. @@ -217,17 +215,14 @@ class MentionedDocumentInfo(BaseModel): id: int title: str = Field(..., min_length=1, max_length=500) - document_type: str = Field(..., min_length=1, max_length=100) + document_type: str | None = Field(default=None, min_length=1, max_length=100) kind: Literal["doc", "folder", "connector"] = Field( default="doc", description=( "Discriminator for the chip's referent: ``doc`` is a " "knowledge-base ``Document`` row, ``folder`` is a " "knowledge-base ``Folder`` row, and ``connector`` is a " - "concrete connected account. Folders carry the sentinel " - "``document_type='FOLDER'`` to keep the frontend dedup key " - "``(kind:document_type:id)`` from colliding doc and folder " - "ids that happen to share an integer value." + "concrete connected account." ), ) connector_type: str | None = Field(default=None, max_length=100) diff --git a/surfsense_backend/app/tasks/chat/persistence.py b/surfsense_backend/app/tasks/chat/persistence.py index 07266cf69..9d100c13c 100644 --- a/surfsense_backend/app/tasks/chat/persistence.py +++ b/surfsense_backend/app/tasks/chat/persistence.py @@ -109,7 +109,7 @@ def _build_user_content( [{"type": "text", "text": "..."}, {"type": "image", "image": "data:..."}, {"type": "mentioned-documents", "documents": [{"id": int, - "title": str, "document_type": str, "kind": "doc" | "folder"}, + "title": str, "kind": "doc" | "folder" | "connector", ...}, ...]}] The companion reader is @@ -117,8 +117,8 @@ def _build_user_content( which expects exactly this shape — keep them in sync. ``mentioned_documents``: optional list of mention chip dicts. Each - dict may include a ``kind`` discriminator (``"doc"`` or ``"folder"``) - so the persisted ContentPart round-trips folder chips on reload. + dict may include a ``kind`` discriminator so the persisted + ContentPart round-trips folder and connector chips on reload. When ``kind`` is missing we default to ``"doc"`` so legacy clients that haven't migrated to the union schema still persist correctly. """ @@ -134,18 +134,23 @@ def _build_user_content( doc_id = doc.get("id") title = doc.get("title") document_type = doc.get("document_type") - if doc_id is None or title is None or document_type is None: - continue kind_raw = doc.get("kind", "doc") kind = kind_raw if kind_raw in ("doc", "folder", "connector") else "doc" + if doc_id is None or title is None: + continue + if kind == "doc" and document_type is None: + continue item = { "id": doc_id, "title": str(title), - "document_type": str(document_type), "kind": kind, } + if document_type is not None: + item["document_type"] = str(document_type) if kind == "connector": - connector_type = doc.get("connector_type") or document_type + connector_type = doc.get("connector_type") + if connector_type is None: + continue account_name = doc.get("account_name") or title item["connector_type"] = str(connector_type) item["account_name"] = str(account_name) diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 8d1f5da46..6cd95a79c 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -73,6 +73,7 @@ import { convertToThreadMessage, reconcileInterruptedAssistantMessages, } from "@/lib/chat/message-utils"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { isPodcastGenerating, looksLikePodcastRequest, @@ -206,7 +207,7 @@ function pairBundleToolCallIds( const MentionedDocumentInfoSchema = z.object({ id: z.number(), title: z.string(), - document_type: z.string(), + document_type: z.string().optional(), kind: z .union([z.literal("doc"), z.literal("folder"), z.literal("connector")]) .optional() @@ -234,9 +235,8 @@ function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] { return { id: doc.id, title: doc.title, - document_type: doc.document_type, kind: "connector", - connector_type: doc.connector_type ?? doc.document_type, + connector_type: doc.connector_type ?? doc.document_type ?? "UNKNOWN", account_name: doc.account_name ?? doc.title, }; } @@ -244,14 +244,13 @@ function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] { return { id: doc.id, title: doc.title, - document_type: "FOLDER", kind: "folder", }; } return { id: doc.id, title: doc.title, - document_type: doc.document_type, + document_type: doc.document_type ?? "UNKNOWN", kind: "doc", }; }); @@ -957,15 +956,13 @@ export default function NewChatPage() { }); // Collect unique mention chips for display & persistence. - // Dedup key is ``kind:document_type:id`` so a folder and a - // doc with the same integer id never collapse into one - // entry. The ``kind`` field is forwarded to the backend + // The ``kind`` field is forwarded to the backend // so the persisted ``mentioned-documents`` content part // can render the correct chip type on reload. const allMentionedDocs: MentionedDocumentInfo[] = []; const seenDocKeys = new Set(); for (const doc of mentionedDocuments) { - const key = `${doc.kind}:${doc.document_type}:${doc.id}`; + const key = getMentionDocKey(doc); if (seenDocKeys.has(key)) continue; seenDocKeys.add(key); allMentionedDocs.push(doc); diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index 9efd2b7fe..25d1e397a 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -3,13 +3,6 @@ import { atom } from "jotai"; import type { Document } from "@/contracts/types/document.types"; -/** - * Sentinel ``document_type`` used for folder mention chips so the - * dedup key (`kind:document_type:id`) never collides a document with a - * folder that happens to share an integer id. - */ -export const FOLDER_MENTION_DOCUMENT_TYPE = "FOLDER"; - /** * Display metadata for a single ``@``-mention chip. * @@ -27,13 +20,11 @@ export type MentionedDocumentInfo = | { id: number; title: string; - document_type: typeof FOLDER_MENTION_DOCUMENT_TYPE; kind: "folder"; } | { id: number; title: string; - document_type: string; kind: "connector"; connector_type: string; account_name: string; @@ -51,8 +42,7 @@ type LegacyDocMention = Pick; * Normalize an arbitrary chip-like input into the discriminated * ``MentionedDocumentInfo`` shape. Existing call sites that only have * ``{id, title, document_type}`` flow through here so they don't have - * to thread ``kind`` everywhere — the helper defaults to ``"doc"`` and - * rewrites the document type for folders. + * to thread ``kind`` everywhere — the helper defaults to ``"doc"``. */ export function toMentionedDocumentInfo( input: LegacyDocMention | MentionedDocumentInfo @@ -78,31 +68,32 @@ export function makeFolderMention(input: { id: number; name: string }): Mentione return { id: input.id, title: input.name, - document_type: FOLDER_MENTION_DOCUMENT_TYPE, kind: "folder", }; } /** - * Atom to store the full mention objects (documents + folders) attached - * via @-mention chips in the current chat composer. Persists across - * component remounts. + * Atom to store the full context objects attached via @-mention chips in + * the current chat composer. Persists across component remounts. */ export const mentionedDocumentsAtom = atom([]); /** * Derived read-only atom that maps deduplicated mention chips into - * backend payload fields. Doc chips split by ``document_type`` exactly - * like before; folder chips are projected into a separate - * ``folder_ids`` bucket so the route can forward - * ``mentioned_folder_ids`` to the agent without the priority middleware - * conflating them with hybrid-search ids. + * backend payload fields. Each mention kind maps to its own explicit + * payload bucket so non-document context never has to masquerade as a + * document type. */ export const mentionedDocumentIdsAtom = atom((get) => { const allMentions = get(mentionedDocumentsAtom); const seen = new Set(); const deduped = allMentions.filter((m) => { - const key = `${m.kind}:${m.document_type}:${m.id}`; + const key = + m.kind === "doc" + ? `doc:${m.document_type}:${m.id}` + : m.kind === "connector" + ? `connector:${m.connector_type}:${m.id}` + : `folder:${m.id}`; if (seen.has(key)) return false; seen.add(key); return true; @@ -120,7 +111,6 @@ export const mentionedDocumentIdsAtom = atom((get) => { connectors: connectors.map((c) => ({ id: c.id, title: c.title, - document_type: c.document_type, kind: c.kind, connector_type: c.connector_type, account_name: c.account_name, diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index b93ea253d..c0d9d9212 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -20,7 +20,6 @@ import { useMemo, useRef, } from "react"; -import { FOLDER_MENTION_DOCUMENT_TYPE } from "@/atoms/chat/mentioned-documents.atom"; import { Button } from "@/components/ui/button"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document } from "@/contracts/types/document.types"; @@ -40,8 +39,6 @@ export interface MentionedDocument { /** * Input shape for inserting a chip. ``kind`` defaults to ``"doc"``. - * Folder chips default ``document_type`` to ``FOLDER_MENTION_DOCUMENT_TYPE`` - * so the dedup key never collides with a doc chip sharing the same id. */ export type MentionChipInput = { id: number; @@ -78,7 +75,12 @@ export interface InlineMentionEditorRef { doc: Pick, options?: { removeTriggerText?: boolean } ) => void; - removeDocumentChip: (docId: number, docType?: string) => void; + removeDocumentChip: ( + docId: number, + docType?: string, + kind?: MentionKind, + connectorType?: string + ) => void; setDocumentChipStatus: ( docId: number, docType: string | undefined, @@ -95,7 +97,7 @@ interface InlineMentionEditorProps { onActionClose?: () => void; onSubmit?: () => void; onChange?: (text: string, docs: MentionedDocument[]) => void; - onDocumentRemove?: (docId: number, docType?: string) => void; + onDocumentRemove?: (docId: number, docType?: string, kind?: MentionKind, connectorType?: string) => void; onKeyDown?: (e: React.KeyboardEvent) => void; disabled?: boolean; className?: string; @@ -135,7 +137,12 @@ const EMPTY_VALUE: ComposerValue = [{ type: "p", children: [{ text: "" }] }]; * the X button and Backspace go through the same call site. */ type MentionEditorContextValue = { - removeChip: (docId: number, docType: string | undefined) => void; + removeChip: ( + docId: number, + docType: string | undefined, + kind: MentionKind | undefined, + connectorType: string | undefined + ) => void; }; const MentionEditorContext = createContext(null); @@ -181,7 +188,12 @@ const MentionElement: FC> = ({ onMouseDown={(e) => e.preventDefault()} onClick={(e) => { e.stopPropagation(); - ctx.removeChip(element.id, element.document_type); + ctx.removeChip( + element.id, + element.document_type, + element.kind, + element.connector_type + ); }} className="absolute inset-0 size-3 rounded-sm p-0 opacity-0 transition-opacity hover:bg-transparent hover:text-primary focus-visible:opacity-100 focus-visible:outline-none focus-visible:ring-0 group-hover:opacity-100 [&_svg]:size-3" > @@ -456,18 +468,11 @@ export const InlineMentionEditor = forwardRef { + (docId: number, docType?: string, kind?: MentionKind, connectorType?: string) => { const match = (n: unknown) => { if (!n || typeof n !== "object" || !("type" in n)) return false; const node = n as MentionElementNode; if (node.type !== MENTION_TYPE) return false; if (node.id !== docId) return false; + if (kind) { + return ( + getMentionDocKey({ + id: node.id, + kind: node.kind ?? "doc", + document_type: node.document_type, + connector_type: node.connector_type, + }) === + getMentionDocKey({ + id: docId, + kind, + document_type: docType, + connector_type: connectorType, + }) + ); + } return (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN"); }; @@ -554,9 +575,14 @@ export const InlineMentionEditor = forwardRef { - removeDocumentChip(docId, docType); - onDocumentRemove?.(docId, docType); + ( + docId: number, + docType: string | undefined, + kind: MentionKind | undefined, + connectorType: string | undefined + ) => { + removeDocumentChip(docId, docType, kind, connectorType); + onDocumentRemove?.(docId, docType, kind, connectorType); }, [onDocumentRemove, removeDocumentChip] ); @@ -679,7 +705,7 @@ export const InlineMentionEditor = forwardRef { } } return docs.map((d) => { - const documentType = d.document_type ?? "UNKNOWN"; if (d.kind === "connector") { return { id: d.id, title: d.title, - document_type: documentType, kind: "connector", - connector_type: d.connector_type ?? documentType, + connector_type: d.connector_type ?? "UNKNOWN", account_name: d.account_name ?? d.title, }; } @@ -559,17 +556,13 @@ const Composer: FC = () => { return { id: d.id, title: d.title, - document_type: FOLDER_MENTION_DOCUMENT_TYPE, kind: "folder", }; } return { id: d.id, title: d.title, - // Atom requires a string; ``"UNKNOWN"`` matches the - // sentinel ``getMentionDocKey`` and the editor's - // match predicates use. - document_type: documentType, + document_type: d.document_type ?? "UNKNOWN", kind: "doc", }; }); @@ -760,13 +753,14 @@ const Composer: FC = () => { ]); const handleDocumentRemove = useCallback( - (docId: number, docType?: string) => { + (docId: number, docType?: string, kind?: "doc" | "folder" | "connector", connectorType?: string) => { setMentionedDocuments((prev) => { - if (!docType) { - // Fallback when chip type is unavailable. - return prev.filter((doc) => doc.id !== docId); - } - const removedKey = getMentionDocKey({ id: docId, document_type: docType }); + const removedKey = getMentionDocKey({ + id: docId, + document_type: docType, + kind, + connector_type: connectorType, + }); return prev.filter((doc) => getMentionDocKey(doc) !== removedKey); }); }, @@ -810,7 +804,12 @@ const Composer: FC = () => { for (const [key, doc] of prevDocsMap) { if (!nextDocsMap.has(key)) { - editor.removeDocumentChip(doc.id, doc.document_type); + editor.removeDocumentChip( + doc.id, + doc.kind === "doc" ? doc.document_type : undefined, + doc.kind, + doc.kind === "connector" ? doc.connector_type : undefined + ); } } diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx index 3e6dc829a..b30db5f69 100644 --- a/surfsense_web/components/assistant-ui/user-message.tsx +++ b/surfsense_web/components/assistant-ui/user-message.tsx @@ -104,7 +104,7 @@ const UserTextPart: FC = () => { const icon = isFolder ? ( ) : isConnector ? ( - getConnectorIcon(segment.doc.connector_type ?? segment.doc.document_type, "size-3.5") ?? ( + getConnectorIcon(segment.doc.connector_type, "size-3.5") ?? ( ) ) : ( diff --git a/surfsense_web/components/documents/FolderTreeView.tsx b/surfsense_web/components/documents/FolderTreeView.tsx index fb1030028..7c076e99a 100644 --- a/surfsense_web/components/documents/FolderTreeView.tsx +++ b/surfsense_web/components/documents/FolderTreeView.tsx @@ -190,7 +190,6 @@ export function FolderTreeView({ for (const f of folders) { const folderMentionKey = getMentionDocKey({ id: f.id, - document_type: "FOLDER", kind: "folder", }); states[f.id] = mentionedDocKeys.has(folderMentionKey) ? "all" : "none"; diff --git a/surfsense_web/components/new-chat/document-mention-picker.tsx b/surfsense_web/components/new-chat/document-mention-picker.tsx index f8a84c51b..c424ae6c3 100644 --- a/surfsense_web/components/new-chat/document-mention-picker.tsx +++ b/surfsense_web/components/new-chat/document-mention-picker.tsx @@ -20,10 +20,7 @@ import { useState, } from "react"; import type * as React from "react"; -import { - FOLDER_MENTION_DOCUMENT_TYPE, - type MentionedDocumentInfo, -} from "@/atoms/chat/mentioned-documents.atom"; +import type { MentionedDocumentInfo } from "@/atoms/chat/mentioned-documents.atom"; import { useAtomValue } from "jotai"; import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms"; import { @@ -78,6 +75,10 @@ type ResourceNodeValue = | { kind: "view"; view: BrowseView } | { kind: "mention"; mention: MentionedDocumentInfo }; +function isConnectorActive(connector: SearchSourceConnector) { + return connector.is_active !== false; +} + function useDebounced(value: T, delay = DEBOUNCE_MS) { const [debounced, setDebounced] = useState(value); const timeoutRef = useRef | undefined>(undefined); @@ -115,22 +116,24 @@ function makeDocMention(doc: Pick): }; } -function makeFolderMention(folder: { id: number; title: string }): MentionedDocumentInfo { +function makeFolderMention( + folder: { id: number; title: string } +): Extract { return { id: folder.id, title: folder.title, - document_type: FOLDER_MENTION_DOCUMENT_TYPE, kind: "folder", }; } -function makeConnectorMention(connector: SearchSourceConnector): MentionedDocumentInfo { +function makeConnectorMention( + connector: SearchSourceConnector +): Extract { const accountName = getConnectorDisplayName(connector.name); const connectorTitle = titleForConnectorType(connector.connector_type); return { id: connector.id, title: `${connectorTitle}: ${accountName}`, - document_type: connector.connector_type, kind: "connector", connector_type: connector.connector_type, account_name: accountName, @@ -140,8 +143,8 @@ function makeConnectorMention(connector: SearchSourceConnector): MentionedDocume function mentionMatchesSearch(mention: MentionedDocumentInfo, searchLower: string) { return [ mention.title, - mention.document_type, mention.kind, + mention.kind === "doc" ? mention.document_type : "", mention.kind === "connector" ? mention.connector_type : "", mention.kind === "connector" ? mention.account_name : "", ].some((value) => value.toLowerCase().includes(searchLower)); @@ -171,6 +174,7 @@ export const DocumentMentionPicker = forwardRef< const [zeroFolders] = useZeroQuery(queries.folders.bySpace({ searchSpaceId })); const { data: connectors = [], isLoading: isConnectorsLoading } = useAtomValue(connectorsAtom); + const activeConnectors = useMemo(() => connectors.filter(isConnectorActive), [connectors]); const paginationScopeKey = useMemo( () => `${searchSpaceId}:${debouncedSearch}`, [searchSpaceId, debouncedSearch] @@ -307,8 +311,8 @@ export const DocumentMentionPicker = forwardRef< }, [zeroFolders, debouncedSearch, deferredSearch, isSingleCharSearch, hasSearch]); const connectorMentions = useMemo( - () => connectors.filter((c) => c.is_active).map(makeConnectorMention), - [connectors] + () => activeConnectors.map(makeConnectorMention), + [activeConnectors] ); const selectedKeys = useMemo( @@ -345,16 +349,16 @@ export const DocumentMentionPicker = forwardRef< { id: "connectors", label: "Connectors", - subtitle: connectors.length + subtitle: activeConnectors.length ? "Choose the exact account for tool use" : "No connected accounts yet", icon: , type: "branch", - disabled: connectors.length === 0, + disabled: activeConnectors.length === 0, value: { kind: "view", view: { kind: "connectors" } }, }, ], - [connectors.length] + [activeConnectors.length] ); const searchNodes = useMemo[]>(() => { @@ -385,7 +389,7 @@ export const DocumentMentionPicker = forwardRef< id: getMentionDocKey(mention), label: mention.title, subtitle: "Connector account", - icon: getConnectorIcon(mention.document_type, "size-4") ?? , + icon: getConnectorIcon(mention.connector_type, "size-4") ?? , type: "item" as const, disabled: selectedKeys.has(getMentionDocKey(mention)), value: { kind: "mention" as const, mention }, @@ -404,7 +408,7 @@ export const DocumentMentionPicker = forwardRef< const connectorTypeEntries = useMemo(() => { const byType = new Map(); - for (const connector of connectors.filter((c) => c.is_active)) { + for (const connector of activeConnectors) { const list = byType.get(connector.connector_type) ?? []; list.push(connector); byType.set(connector.connector_type, list); @@ -412,7 +416,7 @@ export const DocumentMentionPicker = forwardRef< return Array.from(byType.entries()).sort(([a], [b]) => titleForConnectorType(a).localeCompare(titleForConnectorType(b)) ); - }, [connectors]); + }, [activeConnectors]); const browseNodes = useMemo[]>(() => { if (view.kind === "root") return rootNodes; @@ -469,8 +473,8 @@ export const DocumentMentionPicker = forwardRef< }, })); } - return connectors - .filter((connector) => connector.is_active && connector.connector_type === view.connectorType) + return activeConnectors + .filter((connector) => connector.connector_type === view.connectorType) .map((connector) => { const mention = makeConnectorMention(connector); return { @@ -484,7 +488,7 @@ export const DocumentMentionPicker = forwardRef< }; }); }, [ - connectors, + activeConnectors, connectorTypeEntries, folderMentions, rootNodes, diff --git a/surfsense_web/lib/chat/mention-doc-key.ts b/surfsense_web/lib/chat/mention-doc-key.ts index 114faace5..87676dbd6 100644 --- a/surfsense_web/lib/chat/mention-doc-key.ts +++ b/surfsense_web/lib/chat/mention-doc-key.ts @@ -1,18 +1,20 @@ type MentionKeyInput = { id: number; document_type?: string | null; + connector_type?: string | null; kind?: "doc" | "folder" | "connector"; }; /** * Build a stable dedup key for a mention chip. * - * The ``kind:document_type:id`` shape prevents a document and a folder - * with the same integer id from colliding in the chip array (folders - * use the ``FOLDER`` sentinel ``document_type``; the ``kind`` prefix - * is the belt-and-braces guard). + * Each mention kind keys off its real identity fields: + * docs by document type, folders by folder id, and connectors by + * connector type + account id. */ export function getMentionDocKey(doc: MentionKeyInput): string { const kind = doc.kind ?? "doc"; - return `${kind}:${doc.document_type ?? "UNKNOWN"}:${doc.id}`; + if (kind === "folder") return `folder:${doc.id}`; + if (kind === "connector") return `connector:${doc.connector_type ?? "UNKNOWN"}:${doc.id}`; + return `doc:${doc.document_type ?? "UNKNOWN"}:${doc.id}`; } diff --git a/surfsense_web/lib/connector-telemetry.ts b/surfsense_web/lib/connector-telemetry.ts index ef1b3de32..eeccea1e8 100644 --- a/surfsense_web/lib/connector-telemetry.ts +++ b/surfsense_web/lib/connector-telemetry.ts @@ -1,22 +1,17 @@ -import { EnumConnectorName } from "@/contracts/enums/connector"; -import type { SearchSourceConnector } from "@/contracts/types/connector.types"; import { - OAUTH_CONNECTORS, COMPOSIO_CONNECTORS, CRAWLERS, + OAUTH_CONNECTORS, OTHER_CONNECTORS, } from "@/components/assistant-ui/connector-popup/constants/connector-constants"; +import { EnumConnectorName } from "@/contracts/enums/connector"; +import type { SearchSourceConnector } from "@/contracts/types/connector.types"; // ============================================================================= // Connector Telemetry Types & Registry // ============================================================================= -export type ConnectorTelemetryGroup = - | "oauth" - | "composio" - | "crawler" - | "other" - | "unknown"; +export type ConnectorTelemetryGroup = "oauth" | "composio" | "crawler" | "other" | "unknown"; export interface ConnectorTelemetryMeta { connector_type: string; @@ -31,10 +26,11 @@ export interface ConnectorTelemetryMeta { * picked up here, so adding a new integration does NOT require touching * `lib/posthog/events.ts` or per-connector tracking code. */ -const CONNECTOR_TELEMETRY_REGISTRY: ReadonlyMap< - string, - ConnectorTelemetryMeta -> = (() => { +let connectorTelemetryRegistry: ReadonlyMap | undefined; + +function getConnectorTelemetryRegistry(): ReadonlyMap { + if (connectorTelemetryRegistry) return connectorTelemetryRegistry; + const map = new Map(); for (const c of OAUTH_CONNECTORS) { @@ -70,18 +66,17 @@ const CONNECTOR_TELEMETRY_REGISTRY: ReadonlyMap< }); } - return map; -})(); + connectorTelemetryRegistry = map; + return connectorTelemetryRegistry; +} /** * Returns telemetry metadata for a connector_type, or a minimal "unknown" * record so tracking never no-ops for connectors that exist in the backend * but were forgotten in the UI registry. */ -export function getConnectorTelemetryMeta( - connectorType: string, -): ConnectorTelemetryMeta { - const hit = CONNECTOR_TELEMETRY_REGISTRY.get(connectorType); +export function getConnectorTelemetryMeta(connectorType: string): ConnectorTelemetryMeta { + const hit = getConnectorTelemetryRegistry().get(connectorType); if (hit) return hit; return { @@ -101,34 +96,20 @@ export function getConnectorTelemetryMeta( * These are used for connectors that were NOT created via MCP OAuth. */ const LEGACY_REAUTH_ENDPOINTS: Partial> = { - [EnumConnectorName.LINEAR_CONNECTOR]: - "/api/v1/auth/linear/connector/reauth", - [EnumConnectorName.JIRA_CONNECTOR]: - "/api/v1/auth/jira/connector/reauth", - [EnumConnectorName.NOTION_CONNECTOR]: - "/api/v1/auth/notion/connector/reauth", - [EnumConnectorName.GOOGLE_DRIVE_CONNECTOR]: - "/api/v1/auth/google/drive/connector/reauth", - [EnumConnectorName.GOOGLE_GMAIL_CONNECTOR]: - "/api/v1/auth/google/gmail/connector/reauth", - [EnumConnectorName.GOOGLE_CALENDAR_CONNECTOR]: - "/api/v1/auth/google/calendar/connector/reauth", - [EnumConnectorName.COMPOSIO_GOOGLE_DRIVE_CONNECTOR]: - "/api/v1/auth/composio/connector/reauth", - [EnumConnectorName.COMPOSIO_GMAIL_CONNECTOR]: - "/api/v1/auth/composio/connector/reauth", - [EnumConnectorName.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR]: - "/api/v1/auth/composio/connector/reauth", - [EnumConnectorName.ONEDRIVE_CONNECTOR]: - "/api/v1/auth/onedrive/connector/reauth", - [EnumConnectorName.DROPBOX_CONNECTOR]: - "/api/v1/auth/dropbox/connector/reauth", - [EnumConnectorName.CONFLUENCE_CONNECTOR]: - "/api/v1/auth/confluence/connector/reauth", - [EnumConnectorName.TEAMS_CONNECTOR]: - "/api/v1/auth/teams/connector/reauth", - [EnumConnectorName.DISCORD_CONNECTOR]: - "/api/v1/auth/discord/connector/reauth", + [EnumConnectorName.LINEAR_CONNECTOR]: "/api/v1/auth/linear/connector/reauth", + [EnumConnectorName.JIRA_CONNECTOR]: "/api/v1/auth/jira/connector/reauth", + [EnumConnectorName.NOTION_CONNECTOR]: "/api/v1/auth/notion/connector/reauth", + [EnumConnectorName.GOOGLE_DRIVE_CONNECTOR]: "/api/v1/auth/google/drive/connector/reauth", + [EnumConnectorName.GOOGLE_GMAIL_CONNECTOR]: "/api/v1/auth/google/gmail/connector/reauth", + [EnumConnectorName.GOOGLE_CALENDAR_CONNECTOR]: "/api/v1/auth/google/calendar/connector/reauth", + [EnumConnectorName.COMPOSIO_GOOGLE_DRIVE_CONNECTOR]: "/api/v1/auth/composio/connector/reauth", + [EnumConnectorName.COMPOSIO_GMAIL_CONNECTOR]: "/api/v1/auth/composio/connector/reauth", + [EnumConnectorName.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR]: "/api/v1/auth/composio/connector/reauth", + [EnumConnectorName.ONEDRIVE_CONNECTOR]: "/api/v1/auth/onedrive/connector/reauth", + [EnumConnectorName.DROPBOX_CONNECTOR]: "/api/v1/auth/dropbox/connector/reauth", + [EnumConnectorName.CONFLUENCE_CONNECTOR]: "/api/v1/auth/confluence/connector/reauth", + [EnumConnectorName.TEAMS_CONNECTOR]: "/api/v1/auth/teams/connector/reauth", + [EnumConnectorName.DISCORD_CONNECTOR]: "/api/v1/auth/discord/connector/reauth", }; /** @@ -138,9 +119,7 @@ const LEGACY_REAUTH_ENDPOINTS: Partial> = { * the URL from the service key. Legacy OAuth connectors fall back to the * static ``LEGACY_REAUTH_ENDPOINTS`` map. */ -export function getReauthEndpoint( - connector: SearchSourceConnector, -): string | undefined { +export function getReauthEndpoint(connector: SearchSourceConnector): string | undefined { const mcpService = connector.config?.mcp_service as string | undefined; if (mcpService) { return `/api/v1/auth/mcp/${mcpService}/connector/reauth`;