even better

2 files changed, 152 insertions, 42 deletions

diff --git a/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx b/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
index af689f243..91a7adf24 100644
--- a/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
+++ b/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
@@ -16,12 +16,14 @@ import { v4 as uuidv4 } from 'uuid';
 import { ClientUtils, OmitKeys } from '../../../../../ClientUtils';
 import { Doc, DocListCast, Opt } from '../../../../../fields/Doc';
 import { DocData, DocLayout, DocViews } from '../../../../../fields/DocSymbols';
+import { Id } from '../../../../../fields/FieldSymbols';
 import { RichTextField } from '../../../../../fields/RichTextField';
 import { ScriptField } from '../../../../../fields/ScriptField';
 import { CsvCast, DocCast, NumCast, PDFCast, RTFCast, StrCast, VideoCast, AudioCast } from '../../../../../fields/Types';
 import { DocUtils } from '../../../../documents/DocUtils';
 import { CollectionViewType, DocumentType } from '../../../../documents/DocumentTypes';
 import { Docs, DocumentOptions } from '../../../../documents/Documents';
+import { DocServer } from '../../../../DocServer';
 import { DocumentManager } from '../../../../util/DocumentManager';
 import { ImageUtils } from '../../../../util/Import & Export/ImageUtils';
 import { LinkManager } from '../../../../util/LinkManager';
@@ -643,8 +645,9 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             });
 
             // Get the simplified chunk using the document manager
-            const { foundChunk, doc } = this.docManager.getSimplifiedChunkById(chunkId);
+            const { foundChunk, doc, dataDoc } = this.docManager.getSimplifiedChunkById(chunkId);
             console.log('doc: ', doc);
+            console.log('dataDoc: ', dataDoc);
             if (!foundChunk) {
                 if (doc) {
                     console.warn(`Chunk not found in document, ${doc.id}, for chunk ID: ${chunkId}`);
@@ -673,7 +676,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     DocumentManager.Instance.showDocument(doc, { openLocation: OpenWhere.addRight }, () => {});
                     return;
                 }
-                this.handleOtherChunkTypes(foundChunk, citation, doc);
+                this.handleOtherChunkTypes(foundChunk, citation, doc, dataDoc);
                 // Show the chunk text in citation popup
                 let chunkText = citation.direct_text || 'Text content not available';
                 this.showCitationPopup(chunkText);
@@ -795,7 +798,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
      * @param citation The citation object.
      * @param doc The document containing the chunk.
      */
-    handleOtherChunkTypes = (foundChunk: SimplifiedChunk, citation: Citation, doc: Doc) => {
+    handleOtherChunkTypes = (foundChunk: SimplifiedChunk, citation: Citation, doc: Doc, dataDoc?: Doc) => {
         switch (foundChunk.chunkType) {
             case CHUNK_TYPE.IMAGE:
             case CHUNK_TYPE.TABLE:
@@ -838,45 +841,10 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 this.startCitationPopupTimer();
 
                 // Check if the document is a PDF (has a PDF viewer component)
-                const isPDF = PDFCast(doc.data) !== null || doc.type === DocumentType.PDF;
+                const isPDF = PDFCast(dataDoc!.data) !== null || dataDoc!.type === DocumentType.PDF;
 
-                DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
-                    // Add a delay to ensure document is fully loaded and rendered
-                    setTimeout(() => {
-                        try {
-                            // Safety check: ensure the document has views
-                            if (!doc[DocViews] || doc[DocViews].size === 0) {
-                                console.warn('Document views not available yet, retrying...');
-                                this.retryPdfSearch(doc, citation, foundChunk, isPDF, 1);
-                                return;
-                            }
-
-                            const views = Array.from(doc[DocViews]);
-                            if (!views.length) {
-                                console.warn('No document views found, retrying...');
-                                this.retryPdfSearch(doc, citation, foundChunk, isPDF, 1);
-                                return;
-                            }
-
-                            const firstView = views[0] as DocumentView;
-                            if (!firstView || !firstView.ComponentView) {
-                                console.warn('Component view not available yet, retrying...');
-                                this.retryPdfSearch(doc, citation, foundChunk, isPDF, 1);
-                                return;
-                            }
-
-                            const pdfComponent = firstView.ComponentView as PDFBox;
-                            if (isPDF && pdfComponent && citation.direct_text) {
-                                // Use our helper to ensure fuzzy search is enabled and execute the search
-                                this.ensureFuzzySearchAndExecute(pdfComponent, citation.direct_text.trim(), foundChunk.startPage);
-                            }
-                        } catch (error) {
-                            console.error('Error accessing PDF component:', error);
-                            // Retry with exponential backoff
-                            this.retryPdfSearch(doc, citation, foundChunk, isPDF, 1);
-                        }
-                    }, 500); // Initial delay before first attempt
-                });
+                // First ensure document is fully visible before trying to access its views
+                this.ensureDocumentIsVisible(dataDoc!, isPDF, citation, foundChunk, doc);
                 break;
             case CHUNK_TYPE.CSV:
             case CHUNK_TYPE.URL:
@@ -891,6 +859,144 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     };
 
     /**
+     * Ensures a document is fully visible and rendered before performing actions on it
+     * @param doc The document to ensure is visible
+     * @param isPDF Whether this is a PDF document
+     * @param citation The citation information
+     * @param foundChunk The chunk information
+     * @param doc The document to ensure is visible
+     */
+    ensureDocumentIsVisible = (doc: Doc, isPDF: boolean, citation: Citation, foundChunk: SimplifiedChunk, layoutDoc?: Doc) => {
+        try {
+            // First, check if the document already has views and is rendered
+            const hasViews = doc[DocViews] && doc[DocViews].size > 0;
+
+            console.log(`Document ${doc.id}: Current state - hasViews: ${hasViews}, isPDF: ${isPDF}`);
+
+            if (hasViews) {
+                // Document is already rendered, proceed with accessing its view
+                this.processPDFDocumentView(doc, isPDF, citation, foundChunk);
+                return;
+            } else if (layoutDoc) {
+                this.processPDFDocumentView(layoutDoc, isPDF, citation, foundChunk);
+                return;
+            }
+
+            // If document is not rendered yet, show it and wait for it to be ready
+            console.log(`Document ${doc.id} needs to be shown first`);
+
+            // Force document to be rendered by using willZoomCentered: true
+            DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
+                // Wait a bit for the document to be fully rendered (longer than our previous attempts)
+                setTimeout(() => {
+                    // Now manually check if document view exists and is valid
+                    this.verifyAndProcessDocumentView(doc, isPDF, citation, foundChunk, 1);
+                }, 800); // Increased initial delay
+            });
+        } catch (error) {
+            console.error('Error ensuring document visibility:', error);
+            // Show the document anyway as a fallback
+            DocumentManager.Instance.showDocument(doc, { willZoomCentered: true });
+        }
+    };
+
+    /**
+     * Verifies document view exists and processes it, with retries if needed
+     */
+    verifyAndProcessDocumentView = (doc: Doc, isPDF: boolean, citation: Citation, foundChunk: SimplifiedChunk, attempt: number) => {
+        // Diagnostic info
+        console.log(`Verify attempt ${attempt}: Document views for ${doc.id}:`, doc[DocViews] ? `Found ${doc[DocViews].size} views` : 'No views');
+
+        // Double-check document exists in current document system
+        const docExists = DocServer.GetCachedRefField(doc[Id]) !== undefined;
+        if (!docExists) {
+            console.warn(`Document ${doc.id} no longer exists in document system`);
+            return;
+        }
+
+        try {
+            if (!doc[DocViews] || doc[DocViews].size === 0) {
+                if (attempt >= 5) {
+                    console.error(`Maximum verification attempts (${attempt}) reached for document ${doc.id}`);
+
+                    // Last resort: force re-creation of the document view
+                    if (isPDF) {
+                        console.log('Forcing document recreation as last resort');
+                        DocumentManager.Instance.showDocument(doc, {
+                            willZoomCentered: true,
+                        });
+                    }
+                    return;
+                }
+
+                // Let's try explicitly requesting the document be shown again
+                if (attempt > 2) {
+                    console.log(`Attempt ${attempt}: Re-requesting document be shown`);
+                    DocumentManager.Instance.showDocument(doc, {
+                        willZoomCentered: true,
+                        openLocation: attempt % 2 === 0 ? OpenWhere.addRight : undefined,
+                    });
+                }
+
+                // Use exponential backoff for retries
+                const nextDelay = Math.min(2000, 500 * Math.pow(1.5, attempt));
+                console.log(`Scheduling retry ${attempt + 1} in ${nextDelay}ms`);
+
+                setTimeout(() => {
+                    this.verifyAndProcessDocumentView(doc, isPDF, citation, foundChunk, attempt + 1);
+                }, nextDelay);
+                return;
+            }
+
+            this.processPDFDocumentView(doc, isPDF, citation, foundChunk);
+        } catch (error) {
+            console.error(`Error on verification attempt ${attempt}:`, error);
+            if (attempt < 5) {
+                setTimeout(
+                    () => {
+                        this.verifyAndProcessDocumentView(doc, isPDF, citation, foundChunk, attempt + 1);
+                    },
+                    500 * Math.pow(1.5, attempt)
+                );
+            }
+        }
+    };
+
+    /**
+     * Processes a PDF document view once we're sure it exists
+     */
+    processPDFDocumentView = (doc: Doc, isPDF: boolean, citation: Citation, foundChunk: SimplifiedChunk) => {
+        try {
+            const views = Array.from(doc[DocViews] || []);
+            if (!views.length) {
+                console.warn('No document views found in document that should have views');
+                return;
+            }
+
+            const firstView = views[0] as DocumentView;
+            if (!firstView) {
+                console.warn('First view is invalid');
+                return;
+            }
+
+            console.log(`Successfully found document view for ${doc.id}:`, firstView.ComponentView ? `Component: ${firstView.ComponentView.constructor.name}` : 'No component view');
+
+            if (!firstView.ComponentView) {
+                console.warn('Component view not available');
+                return;
+            }
+
+            // For PDF documents, perform fuzzy search
+            if (isPDF && firstView.ComponentView && citation.direct_text) {
+                const pdfComponent = firstView.ComponentView as PDFBox;
+                this.ensureFuzzySearchAndExecute(pdfComponent, citation.direct_text.trim(), foundChunk.startPage);
+            }
+        } catch (error) {
+            console.error('Error processing PDF document view:', error);
+        }
+    };
+
+    /**
      * Creates an annotation highlight on a PDF document for image citations.
      * @param x1 X-coordinate of the top-left corner of the highlight.
      * @param y1 Y-coordinate of the top-left corner of the highlight.
diff --git a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
index 784e90c3c..33eec5972 100644
--- a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
+++ b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
@@ -1007,6 +1007,10 @@ export class AgentDocumentManager {
         return docInfo?.layoutDoc;
     }
 
+    public getDataDocument(docId: string): Doc | undefined {
+        const docInfo = this.documentsById.get(docId);
+        return docInfo?.dataDoc;
+    }
     /**
      * Adds simplified chunks to a document for citation handling
      * @param doc The document to add simplified chunks to
@@ -1074,7 +1078,7 @@ export class AgentDocumentManager {
      */
     @action
     public getSimplifiedChunkById(chunkId: string): any | undefined {
-        return { foundChunk: this.simplifiedChunks.get(chunkId), doc: this.getDocument(this.simplifiedChunks.get(chunkId)?.doc_id || chunkId) };
+        return { foundChunk: this.simplifiedChunks.get(chunkId), doc: this.getDocument(this.simplifiedChunks.get(chunkId)?.doc_id || chunkId), dataDoc: this.getDataDocument(this.simplifiedChunks.get(chunkId)?.doc_id || chunkId) };
     }
 
     /**