From 7a914ce2d9b21c835f9a72de730e755c34359194 Mon Sep 17 00:00:00 2001
From: ValueOn AG <p.motsch@poweron.swiss>
Date: Thu, 4 Jun 2026 23:56:58 +0200
Subject: [PATCH] fix: resolve datasource labels on reload using backend labels

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 src/pages/views/workspace/WorkspaceInput.tsx | 55 +++++++++++++++++---
 src/pages/views/workspace/WorkspacePage.tsx  |  2 +
 src/pages/views/workspace/useWorkspace.ts    | 12 +++++
 3 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/src/pages/views/workspace/WorkspaceInput.tsx b/src/pages/views/workspace/WorkspaceInput.tsx
index 604331a..267b1d0 100644
--- a/src/pages/views/workspace/WorkspaceInput.tsx
+++ b/src/pages/views/workspace/WorkspaceInput.tsx
@@ -9,6 +9,7 @@ import React, {
   useRef,
   useEffect,
   useImperativeHandle,
+  useMemo,
   forwardRef,
 } from 'react';
 import { ProviderMultiSelect } from '../../../components/ProviderSelector';
@@ -69,6 +70,8 @@ interface WorkspaceInputProps {
   workflowId?: string | null;
   loadedAttachedDataSourceIds?: string[];
   loadedAttachedFeatureDataSourceIds?: string[];
+  loadedDataSourceLabels?: Record<string, string>;
+  loadedFeatureDataSourceLabels?: Record<string, string>;
   loadedNonce?: number;
 }
 
@@ -135,6 +138,8 @@ export const WorkspaceInput = forwardRef<WorkspaceInputHandle, WorkspaceInputPro
   workflowId,
   loadedAttachedDataSourceIds,
   loadedAttachedFeatureDataSourceIds,
+  loadedDataSourceLabels,
+  loadedFeatureDataSourceLabels,
   loadedNonce,
 }, ref) {
   const { t } = useLanguage();
@@ -154,6 +159,34 @@ export const WorkspaceInput = forwardRef<WorkspaceInputHandle, WorkspaceInputPro
   const [neutralizeActive, setNeutralizeActive] = useState(false);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const dsLabelCache = useRef<Map<string, string>>(new Map());
+  const dsLabelMap = useMemo(() => {
+    const m = new Map<string, string>();
+    if (loadedDataSourceLabels) {
+      for (const [k, v] of Object.entries(loadedDataSourceLabels)) {
+        if (v) m.set(k, v);
+      }
+    }
+    dsLabelCache.current.forEach((v, k) => { if (!m.has(k)) m.set(k, v); });
+    for (const ds of dataSources) {
+      const lbl = ds.label || ds.path;
+      if (lbl) m.set(ds.id, lbl);
+    }
+    return m;
+  }, [dataSources, loadedDataSourceLabels]);
+
+  const fdsLabelMap = useMemo(() => {
+    const m = new Map<string, string>();
+    if (loadedFeatureDataSourceLabels) {
+      for (const [k, v] of Object.entries(loadedFeatureDataSourceLabels)) {
+        if (v) m.set(k, v);
+      }
+    }
+    for (const fds of featureDataSources) {
+      const lbl = fds.label || fds.tableName;
+      if (lbl) m.set(fds.id, lbl);
+    }
+    return m;
+  }, [featureDataSources, loadedFeatureDataSourceLabels]);
 
   const _appendAttachment = useCallback((item: AttachmentItem) => {
     setAttachments(prev => prev.some(a => a.id === item.id) ? prev : [...prev, item]);
@@ -234,26 +267,32 @@ export const WorkspaceInput = forwardRef<WorkspaceInputHandle, WorkspaceInputPro
   useEffect(() => {
     if (loadedNonce === undefined) return;
     if (_reconciledDsForNonce.current === loadedNonce) return;
-    if (dataSources.length === 0) return;
+    if (dataSources.length === 0 && !loadedDataSourceLabels) return;
     _reconciledDsForNonce.current = loadedNonce;
+    for (const ds of dataSources) {
+      const lbl = ds.label || ds.path;
+      if (lbl) dsLabelCache.current.set(ds.id, lbl);
+    }
     const validIds = new Set(dataSources.map(d => d.id));
+    const labeledIds = new Set(Object.keys(loadedDataSourceLabels || {}));
     setAttachedDataSourceIds(prev => {
-      const filtered = prev.filter(id => validIds.has(id));
+      const filtered = prev.filter(id => validIds.has(id) || labeledIds.has(id));
       return filtered.length === prev.length ? prev : filtered;
     });
-  }, [loadedNonce, dataSources]);
+  }, [loadedNonce, dataSources, loadedDataSourceLabels]);
 
   useEffect(() => {
     if (loadedNonce === undefined) return;
     if (_reconciledFdsForNonce.current === loadedNonce) return;
-    if (featureDataSources.length === 0) return;
+    if (featureDataSources.length === 0 && !loadedFeatureDataSourceLabels) return;
     _reconciledFdsForNonce.current = loadedNonce;
     const validIds = new Set(featureDataSources.map(d => d.id));
+    const labeledIds = new Set(Object.keys(loadedFeatureDataSourceLabels || {}));
     setAttachedFeatureDataSourceIds(prev => {
-      const filtered = prev.filter(id => validIds.has(id));
+      const filtered = prev.filter(id => validIds.has(id) || labeledIds.has(id));
       return filtered.length === prev.length ? prev : filtered;
     });
-  }, [loadedNonce, featureDataSources]);
+  }, [loadedNonce, featureDataSources, loadedFeatureDataSourceLabels]);
 
   const promptBeforeVoiceRef = useRef('');
   const finalizedTextRef = useRef('');
@@ -663,7 +702,7 @@ export const WorkspaceInput = forwardRef<WorkspaceInputHandle, WorkspaceInputPro
                   background: '#e8f5e9', color: '#2e7d32', fontWeight: 500,
                 }}
               >
-                🔗 {ds?.label || ds?.path || dsLabelCache.current.get(dsId) || dsId}
+                🔗 {ds?.label || ds?.path || dsLabelMap.get(dsId) || dsId}
                 <button
                   type="button"
                   onClick={() => _removeAttachedDataSource(dsId)}
@@ -690,7 +729,7 @@ export const WorkspaceInput = forwardRef<WorkspaceInputHandle, WorkspaceInputPro
                 }}
               >
                 <span style={{ display: 'flex', alignItems: 'center', fontSize: 12 }}>{fdsIcon || '\uD83D\uDDC3\uFE0F'}</span>
-                {fds?.label || fdsId} – {fds?.tableName || ''}
+                {fds?.label || fdsLabelMap.get(fdsId) || fdsId} – {fds?.tableName || ''}
                 <button
                   type="button"
                   onClick={() => _toggleFeatureDataSource(fdsId)}
diff --git a/src/pages/views/workspace/WorkspacePage.tsx b/src/pages/views/workspace/WorkspacePage.tsx
index a176f19..c53bace 100644
--- a/src/pages/views/workspace/WorkspacePage.tsx
+++ b/src/pages/views/workspace/WorkspacePage.tsx
@@ -571,6 +571,8 @@ export const WorkspacePage: React.FC<WorkspacePageProps> = ({ persistentInstance
           workflowId={workspace.workflowId}
           loadedAttachedDataSourceIds={workspace.loadedAttachedDataSourceIds}
           loadedAttachedFeatureDataSourceIds={workspace.loadedAttachedFeatureDataSourceIds}
+          loadedDataSourceLabels={workspace.loadedDataSourceLabels}
+          loadedFeatureDataSourceLabels={workspace.loadedFeatureDataSourceLabels}
           loadedNonce={workspace.loadedNonce}
         />
       </main>
diff --git a/src/pages/views/workspace/useWorkspace.ts b/src/pages/views/workspace/useWorkspace.ts
index dfe9c50..c95aae1 100644
--- a/src/pages/views/workspace/useWorkspace.ts
+++ b/src/pages/views/workspace/useWorkspace.ts
@@ -120,6 +120,8 @@ interface UseWorkspaceReturn {
    */
   loadedAttachedDataSourceIds: string[];
   loadedAttachedFeatureDataSourceIds: string[];
+  loadedDataSourceLabels: Record<string, string>;
+  loadedFeatureDataSourceLabels: Record<string, string>;
   loadedNonce: number;
 }
 
@@ -138,6 +140,8 @@ export function useWorkspace(instanceId: string): UseWorkspaceReturn {
   const [dataSourceAccesses, setDataSourceAccesses] = useState<DataSourceAccessEvent[]>([]);
   const [loadedAttachedDataSourceIds, setLoadedAttachedDataSourceIds] = useState<string[]>([]);
   const [loadedAttachedFeatureDataSourceIds, setLoadedAttachedFeatureDataSourceIds] = useState<string[]>([]);
+  const [loadedDataSourceLabels, setLoadedDataSourceLabels] = useState<Record<string, string>>({});
+  const [loadedFeatureDataSourceLabels, setLoadedFeatureDataSourceLabels] = useState<Record<string, string>>({});
   const [loadedNonce, setLoadedNonce] = useState(0);
   const cleanupRef = useRef<(() => void) | null>(null);
 
@@ -192,8 +196,12 @@ export function useWorkspace(instanceId: string): UseWorkspaceReturn {
         const fdsIds: string[] = Array.isArray(res.data.attachedFeatureDataSourceIds)
           ? res.data.attachedFeatureDataSourceIds.map((x: unknown) => String(x))
           : [];
+        const dsLabels: Record<string, string> = res.data.attachedDataSourceLabels || {};
+        const fdsLabels: Record<string, string> = res.data.attachedFeatureDataSourceLabels || {};
         setLoadedAttachedDataSourceIds(dsIds);
         setLoadedAttachedFeatureDataSourceIds(fdsIds);
+        setLoadedDataSourceLabels(dsLabels);
+        setLoadedFeatureDataSourceLabels(fdsLabels);
         setLoadedNonce(n => n + 1);
       })
       .catch(() => {});
@@ -208,6 +216,8 @@ export function useWorkspace(instanceId: string): UseWorkspaceReturn {
     setDataSourceAccesses([]);
     setLoadedAttachedDataSourceIds([]);
     setLoadedAttachedFeatureDataSourceIds([]);
+    setLoadedDataSourceLabels({});
+    setLoadedFeatureDataSourceLabels({});
     setLoadedNonce(n => n + 1);
   }, []);
 
@@ -528,6 +538,8 @@ export function useWorkspace(instanceId: string): UseWorkspaceReturn {
     dataSourceAccesses,
     loadedAttachedDataSourceIds,
     loadedAttachedFeatureDataSourceIds,
+    loadedDataSourceLabels,
+    loadedFeatureDataSourceLabels,
     loadedNonce,
   };
 }