11import React , { useState , useCallback , useEffect , useMemo , useRef } from 'react' ;
22import { createPortal } from 'react-dom' ;
3- import { useSearchParams } from 'react-router-dom' ;
3+ import { useSearchParams , useNavigate } from 'react-router-dom' ;
4+ import { parseQueryForHandover } from '../utils/queryHandover' ;
45import { generateMongoQuery , debugMongoQuery , analyzeQueryResult , inferSchemaRelationships , evaluateWriteResult } from '../services/geminiService' ;
56import { getAzureCosmosAccounts , getDatabasesForAccount , runMongoQuery , getCollectionInfo , clearSystemCache } from '../services/dbService' ;
67import { getSavedQueries , saveQuery , updateSavedQuery , deleteSavedQuery } from '../services/userDataService' ;
@@ -408,7 +409,16 @@ export interface QueryGeneratorPageProps {
408409
409410// --- Main Page Component ---
410411const QueryGeneratorPage : React . FC < QueryGeneratorPageProps > = ( { name, email, onLogout, onNavigateToExplorer, onConnectionChange, embedded = false , preselectedAccountId } ) => {
411- const [ userInput , setUserInput ] = useState < string > ( '' ) ;
412+ const navigate = useNavigate ( ) ;
413+
414+ const [ userInput , setUserInput ] = useState < string > ( ( ) => {
415+ try {
416+ const ws = JSON . parse ( sessionStorage . getItem ( 'qp_workspace' ) ?? 'null' ) ;
417+ const conn = JSON . parse ( sessionStorage . getItem ( 'qp_connection' ) ?? 'null' ) ;
418+ if ( ws && conn && ws . accountId === conn . accountId && ws . databaseName === conn . databaseName ) return ws . userInput ?? '' ;
419+ } catch { /* ignore */ }
420+ return '' ;
421+ } ) ;
412422 const [ isLoading , setIsLoading ] = useState < boolean > ( false ) ;
413423 const [ error , setError ] = useState < string | null > ( null ) ;
414424
@@ -434,16 +444,44 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
434444 const [ quickExploringAccountId , setQuickExploringAccountId ] = useState < string | null > ( null ) ;
435445
436446 // State for AI query generation
437- const [ _queryResult , setQueryResult ] = useState < QueryResultData | null > ( null ) ;
438- const [ querySourceCollection , setQuerySourceCollection ] = useState < string | null > ( null ) ;
439- const [ editableCode , setEditableCode ] = useState < string > ( '' ) ;
447+ const [ _queryResult , setQueryResult ] = useState < QueryResultData | null > ( ( ) => {
448+ try {
449+ const ws = JSON . parse ( sessionStorage . getItem ( 'qp_workspace' ) ?? 'null' ) ;
450+ const conn = JSON . parse ( sessionStorage . getItem ( 'qp_connection' ) ?? 'null' ) ;
451+ if ( ws && conn && ws . accountId === conn . accountId && ws . databaseName === conn . databaseName ) return ws . queryResult ?? null ;
452+ } catch { /* ignore */ }
453+ return null ;
454+ } ) ;
455+ const [ querySourceCollection , setQuerySourceCollection ] = useState < string | null > ( ( ) => {
456+ try {
457+ const ws = JSON . parse ( sessionStorage . getItem ( 'qp_workspace' ) ?? 'null' ) ;
458+ const conn = JSON . parse ( sessionStorage . getItem ( 'qp_connection' ) ?? 'null' ) ;
459+ if ( ws && conn && ws . accountId === conn . accountId && ws . databaseName === conn . databaseName ) return ws . querySourceCollection ?? null ;
460+ } catch { /* ignore */ }
461+ return null ;
462+ } ) ;
463+ const [ editableCode , setEditableCode ] = useState < string > ( ( ) => {
464+ try {
465+ const ws = JSON . parse ( sessionStorage . getItem ( 'qp_workspace' ) ?? 'null' ) ;
466+ const conn = JSON . parse ( sessionStorage . getItem ( 'qp_connection' ) ?? 'null' ) ;
467+ if ( ws && conn && ws . accountId === conn . accountId && ws . databaseName === conn . databaseName ) return ws . editableCode ?? '' ;
468+ } catch { /* ignore */ }
469+ return '' ;
470+ } ) ;
440471 const [ codeHistory , setCodeHistory ] = useState < string [ ] > ( [ ] ) ;
441472 const [ historyIndex , setHistoryIndex ] = useState < number > ( - 1 ) ;
442473 const [ lastSuccessfulPrompt , setLastSuccessfulPrompt ] = useState < string > ( '' ) ;
443474
444475 // State for query execution
445476 const [ isExecuting , setIsExecuting ] = useState < boolean > ( false ) ;
446- const [ executionResult , setExecutionResult ] = useState < any | null > ( null ) ;
477+ const [ executionResult , setExecutionResult ] = useState < any | null > ( ( ) => {
478+ try {
479+ const ws = JSON . parse ( sessionStorage . getItem ( 'qp_workspace' ) ?? 'null' ) ;
480+ const conn = JSON . parse ( sessionStorage . getItem ( 'qp_connection' ) ?? 'null' ) ;
481+ if ( ws && conn && ws . accountId === conn . accountId && ws . databaseName === conn . databaseName ) return ws . executionResult ?? null ;
482+ } catch { /* ignore */ }
483+ return null ;
484+ } ) ;
447485 const [ executionError , setExecutionError ] = useState < string | null > ( null ) ;
448486
449487 // State for intermediate context (multi-step queries)
@@ -526,6 +564,50 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
526564 return azureAccounts . find ( acc => acc . id === connectedResource . accountId ) ?. name ?? 'Unknown Account' ;
527565 } , [ connectedResource , azureAccounts ] ) ;
528566
567+ // Query handover — detect transferable find() queries for "Open in Explorer"
568+ const handover = useMemo (
569+ ( ) => editableCode && ! _queryResult ?. is_write_action ? parseQueryForHandover ( editableCode ) : null ,
570+ [ editableCode , _queryResult ]
571+ ) ;
572+
573+ const handleOpenInExplorer = useCallback ( ( ) => {
574+ if ( ! handover || ! connectedResource || ! connectedDbInfo ) return ;
575+ navigate (
576+ `/data-explorer/${ encodeURIComponent ( connectedResource . accountId ) } ${ encodeURIComponent ( connectedResource . databaseName ) } ,
577+ {
578+ state : {
579+ dbInfo : connectedDbInfo ,
580+ accountName : connectedAccountName ,
581+ availableDbs : accountDatabases ,
582+ availableAccounts : azureAccounts ,
583+ initialCollection : handover . collection ,
584+ initialFilters : handover . filters ,
585+ } ,
586+ }
587+ ) ;
588+ } , [ handover , connectedResource , connectedDbInfo , connectedAccountName , accountDatabases , azureAccounts , navigate ] ) ;
589+
590+ // Persist workspace state so it survives sidebar tab switches
591+ useEffect ( ( ) => {
592+ if ( ! connectedResource ) return ;
593+ try {
594+ const payload : Record < string , unknown > = {
595+ accountId : connectedResource . accountId ,
596+ databaseName : connectedResource . databaseName ,
597+ userInput,
598+ editableCode,
599+ querySourceCollection,
600+ queryResult : _queryResult ,
601+ } ;
602+ // Only persist executionResult if it's under 1 MB to avoid quota errors
603+ if ( executionResult !== null ) {
604+ const resultStr = JSON . stringify ( executionResult ) ;
605+ if ( resultStr . length < 1_000_000 ) payload . executionResult = executionResult ;
606+ }
607+ sessionStorage . setItem ( 'qp_workspace' , JSON . stringify ( payload ) ) ;
608+ } catch { /* ignore */ }
609+ } , [ userInput , editableCode , querySourceCollection , connectedResource , _queryResult , executionResult ] ) ;
610+
529611 const [ isWaitingForAuth , setIsWaitingForAuth ] = useState < boolean > ( false ) ;
530612
531613 const fetchAccounts = useCallback ( async ( ) => {
@@ -636,6 +718,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
636718 setCurrentQueryContextSource ( null ) ;
637719 setRelationships ( null ) ; // Clear relationships
638720 setRelationshipError ( null ) ; // Clear relationship error
721+ sessionStorage . removeItem ( 'qp_workspace' ) ;
639722 } , [ ] ) ;
640723
641724 const handleDisconnect = useCallback ( ( ) => {
@@ -695,7 +778,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
695778 }
696779 } , [ selectedAccountId , connectedResource , azureAccounts , handleDisconnect ] ) ;
697780
698- const handleConnectDatabase = useCallback ( async ( dbInfo : DbInfo ) => {
781+ const handleConnectDatabase = useCallback ( async ( dbInfo : DbInfo , preserveWorkspace = false ) => {
699782 const account = azureAccounts . find ( acc => acc . id === selectedAccountId ) ;
700783 if ( ! account ) return ;
701784
@@ -708,7 +791,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
708791 setConnectedDbInfo ( dbInfo ) ;
709792 sessionStorage . setItem ( 'qp_connection' , JSON . stringify ( { accountId : account . id , databaseName : dbInfo . name , accountName : account . name , collections : dbInfo . collections , availableAccounts : azureAccounts , availableDbs : accountDatabases } ) ) ;
710793 onConnectionChange ?.( account . id , dbInfo . name , account . name , dbInfo . collections , azureAccounts , accountDatabases ) ;
711- clearQueryState ( ) ;
794+ if ( ! preserveWorkspace ) clearQueryState ( ) ;
712795 setIsConnectingToDb ( null ) ;
713796 } , [ selectedAccountId , azureAccounts , clearQueryState , onConnectionChange ] ) ;
714797
@@ -741,9 +824,9 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
741824 ) {
742825 if ( targetDatabaseName ) {
743826 const db = accountDatabases . find ( d => d . name === targetDatabaseName ) ?? accountDatabases [ 0 ] ;
744- handleConnectDatabase ( db ) ;
827+ handleConnectDatabase ( db , true ) ;
745828 } else if ( accountDatabases . length === 1 ) {
746- handleConnectDatabase ( accountDatabases [ 0 ] ) ;
829+ handleConnectDatabase ( accountDatabases [ 0 ] , true ) ;
747830 }
748831 }
749832 } , [ preselectedAccountId , initialConnection , selectedAccountId , accountDatabases , connectedResource , isLoadingDatabases , isConnectingToDb , handleConnectDatabase ] ) ;
@@ -2389,7 +2472,7 @@ const QueryGeneratorPage: React.FC<QueryGeneratorPageProps> = ({ name, email, on
23892472 { ( ! isLoading && ! error && ! isDemoModeForResultsStep && ! isDemoModeForDebugStep && ! isDemoModeForContextActiveStep && ! isDemoModeForRunStep && ! isDemoModeForSaveStep ) && (
23902473 editableCode ? (
23912474 < div style = { { display : 'flex' , flexDirection : 'column' , gap : 12 } } >
2392- < QueryDisplay code = { editableCode } onCodeChange = { setEditableCode } onRunQuery = { handleRunQuery } onSaveQuery = { handleOpenSaveDialog } isExecuting = { isExecuting } historyCount = { codeHistory . length } historyIndex = { historyIndex } onNavigateHistory = { handleNavigateHistory } />
2475+ < QueryDisplay code = { editableCode } onCodeChange = { setEditableCode } onRunQuery = { handleRunQuery } onSaveQuery = { handleOpenSaveDialog } isExecuting = { isExecuting } historyCount = { codeHistory . length } historyIndex = { historyIndex } onNavigateHistory = { handleNavigateHistory } isTransferable = { ! ! handover } onOpenInExplorer = { handleOpenInExplorer } />
23932476 < QueryResult isExecuting = { isExecuting } executionError = { executionError } executionResult = { executionResult } onDebug = { handleDebugQuery } isDebugging = { isDebugging } debuggingResult = { debuggingResult } debugError = { debugError } sourceCollection = { querySourceCollection } onSetIntermediateContext = { handleSetIntermediateContext } intermediateContext = { intermediateContext } onAnalyze = { handleAnalyzeQuery } isAnalyzing = { isAnalyzing } analysisResult = { analysisResult } analysisError = { analysisError } onEvaluateWrite = { handleEvaluateWrite } isEvaluatingWrite = { isEvaluatingWrite } writeEvaluationResult = { writeEvaluationResult } writeEvaluationError = { writeEvaluationError } />
23942477 </ div >
23952478 ) : (
0 commit comments