Feature/embed auth (#276)

* Add embedded mx token flow

* Implement embedded auth flow

* Use embed login flow in embedded mode

Author: ppsreejith

extension/src/content/RPCs/fetchData.ts

@@ -15,6 +15,11 @@ function getCookie(name: string): string | null {
   return cookieValue
 }
 
+export function getMXToken(): string | null {
+  // Get the CSRF token from the cookie
+  return getCookie('mx_jwt')
+}
+
 export async function fetchData<T>(
   url: string,
   method: HttpMethod,

extension/src/content/RPCs/index.ts

@@ -6,7 +6,7 @@ import { captureVisibleTab } from "./rpcCalls"
 import { copyToClipboard } from "./copyToClipboard"
 import { getElementScreenCapture } from "./elementScreenCapture"
 import ripple from "./ripple"
-import { fetchData } from "./fetchData"
+import { fetchData, getMXToken } from "./fetchData"
 import { initWindowListener, RPCPayload } from './initListeners'
 import { addNativeElements, attachEventsListener, attachMutationListener, detachMutationListener, initMutationObserver } from "./mutationObserver"
 import { respondToOtherTab, forwardToTab, getPendingMessage } from "./crossInstanceComms"
@@ -43,6 +43,7 @@ export const rpc = {
     stopRecording,
     attachEventsListener,
     addNativeElements,
+    getMXToken
 }
 
 type RPC = typeof rpc

web/src/app/api/auth.ts

@@ -77,5 +77,14 @@ export default {
       }
       return {}
     }
+  },
+
+  async embedAuth(token: string) {
+    try {
+      const response = await axios.post(`${API_BASE_URL}/embed`, { token });
+      return response.data;
+    } catch (error) {
+      console.error('Error during embed auth:', error);
+    }
   }
-}
+}

web/src/app/rpc.ts

@@ -106,6 +106,7 @@ export const fetchData = (
     log_rpc: true,
   })
 export const queryURL = () => sendMessage('queryURL', [])
+export const getMXToken = () => sendMessage('getMXToken', [])
 export const getMetabaseState = (path: Parameters<typeof get>[1]) =>
   sendMessage('getMetabaseState', [path], { log_rpc: false })
 export const dispatchMetabaseAction = (type: string, payload?: any) =>

web/src/components/common/App.tsx

@@ -149,6 +149,7 @@ const AppLoggedIn = forwardRef((_props, ref) => {
   const analystMode = useSelector((state: RootState) => state.settings.analystMode)
   const drMode = useSelector((state: RootState) => state.settings.drMode)
   const currentAgent: AgentType = analystMode ? AGENTS.EXPLORER : drMode ? AGENTS.SIMPLE : AGENTS.CLASSIC
+  const isEmbedded = getParsedIframeInfo().isEmbedded as unknown === 'true'
 
   const agentIconMap: Record<AgentType, any> = {
     [AGENTS.EXPLORER]: BiSolidMapAlt,
@@ -190,7 +191,7 @@ const AppLoggedIn = forwardRef((_props, ref) => {
   // Update thread id on start
   useEffect(() => {
     dispatch(updateThreadID())
-  }, [])
+  }, []) 
 
   useEffect(() => {
     getBillingInfo().then(billingInfo => {
@@ -214,7 +215,7 @@ const AppLoggedIn = forwardRef((_props, ref) => {
   }, [])
   useEffect(() => {
     const attemptRefresh = () => {
-      if (configs.IS_DEV) {
+      if (configs.IS_DEV || isEmbedded) {
         return
       }
       authModule.refresh().then(({ expired, changed, session_jwt, profile_id, email }) => {
@@ -246,7 +247,6 @@ const AppLoggedIn = forwardRef((_props, ref) => {
 //   const isDevToolsOpen = useSelector((state: RootState) => state.settings.isDevToolsOpen)
   const platformShortcut = getPlatformShortcut()
   const width = getParsedIframeInfo().width
-  const isEmbedded = getParsedIframeInfo().isEmbedded as unknown === 'true'
 
   const clearMessages = () => {
     if (taskInProgress) {

web/src/components/common/Auth.tsx

@@ -2,7 +2,7 @@ import React, { useState, useRef, useEffect } from 'react';
 import { Button, Input, Box, VStack, Image, CloseButton, HStack, Text, Progress } from '@chakra-ui/react';
 import { Select, CreatableSelect } from "chakra-react-select";
 import { login } from '../../state/auth/reducer'
-import { dispatch } from '../../state/dispatch'
+import { dispatch, logoutState, resetState } from '../../state/dispatch'
 import {auth, auth as authModule} from '../../app/api'
 import { useSelector } from 'react-redux';
 import logo from '../../assets/img/logo.svg'
@@ -14,10 +14,11 @@ import { TelemetryToggle } from '../devtools/Settings';
 import { getParsedIframeInfo } from '../../helpers/origin';
 import { toast } from '../../app/toast';
 import { get } from 'lodash';
-import { setMinusxMode } from '../../app/rpc';
+import { getMXToken, setMinusxMode } from '../../app/rpc';
 import { RootState } from '../../state/store';
 import { getApp } from '../../helpers/app';
 import { Markdown } from './Markdown';
+import { setCurrentEmail } from '../../state/settings/reducer';
 
 
 const useAppStore = getApp().useStore()
@@ -177,6 +178,35 @@ const Auth = () => {
   const [discoverySource, setDiscoverySource] = useState("");
   const isOTPMode = authJWT ? true : false
   const helperMessage = useAppStore((state) => state.helperMessage)?.split('---')[1] || "Welcome to MinusX! You can ask us anything related to your data, and our agents will take care of the rest!"
+  const isEmbedded = getParsedIframeInfo().isEmbedded as unknown === 'true'
+  const currentEmail = useSelector((state: RootState) => state.settings.currentEmail)
+
+  useEffect(() => {
+    const checkToken = async () => {
+      const mx_token = await getMXToken()
+      
+      if (mx_token) {
+        try {
+          const embedAuthResult = await authModule.embedAuth(mx_token)
+          const { session_jwt, profile_id, email } = embedAuthResult
+          if (email != currentEmail) {
+            resetState()
+          }
+          dispatch(login({
+            session_jwt,
+            profile_id,
+            email,
+          }))
+          dispatch(setCurrentEmail(email))
+        } catch (error) {
+          console.error('Failed to authenticate embed token:', error)
+        }
+      }
+    }
+    if (isEmbedded) {
+      checkToken()
+    }
+  }, [])
 
 
   useEffect(() => {
@@ -217,6 +247,7 @@ const Auth = () => {
           profile_id,
           email,
       }))
+      dispatch(setCurrentEmail(email))
       if (is_new_user) {
         captureEvent(GLOBAL_EVENTS.user_signup, { email, profile_id, discoverySource })
       } else {
@@ -330,12 +361,13 @@ const Auth = () => {
     borderWidth={1.5} borderLeftColor={"minusxBW.500"}>
       <Image src={logo} alt="MinusX" maxWidth='150px'/>
       <VStack spacing={4} mt={5} position={"relative"}>
+        {isEmbedded && <Text>Logging you in...</Text>}
         <Input
           type="email"
           placeholder="Enter work email ID"
           aria-label="Enter work email ID"
           value={email}
-          disabled={isOTPMode ? true : false}
+          disabled={isEmbedded ? true : isOTPMode ? true : false}
           onChange={(e) => setEmail(e.target.value)}
           // just trigger the handleSignin function when enter is pressed in this input as well
           onKeyUp={(e) => {

web/src/components/devtools/index.tsx

@@ -69,7 +69,7 @@ export const DevToolsBox: React.FC = () => {
     // Check existing dev/production logic
     const isAllowedByEnv = configs.IS_DEV || Monitor.tags?.includes('production')
 
-    if (isEmbedded) {
+    if (isEmbedded && !configs.IS_DEV) {
         if (Monitor.title === 'History' || Monitor.title === 'Memory') {
             return true
         }

web/src/state/settings/reducer.ts

@@ -133,6 +133,7 @@ interface Settings {
   metadataHashes: Record<string, number>
   metadataProcessingCache: Record<number, MetadataProcessingCacheEntry>
   manuallyLimitContext: boolean
+  currentEmail?: string
 }
 
 const initialState: Settings = {
@@ -224,6 +225,9 @@ export const settingsSlice = createSlice({
     setAppRecording: (state, action: PayloadAction<boolean>) => {
       state.isRecording = action.payload
     },
+    setCurrentEmail: (state, action: PayloadAction<string | undefined>) => {
+      state.currentEmail = action.payload
+    },
     setAiRules: (state, action: PayloadAction<string>) => {
       state.aiRules = action.payload
     },
@@ -412,7 +416,7 @@ export const { updateIsLocal, updateUploadLogs,
   setIframeInfo, setConfirmChanges, setDemoMode, setAppRecording, setAiRules,
   applyTableDiff, setSelectedModels, setDRMode, setAnalystMode, setSelectedCatalog, saveCatalog, deleteCatalog, setMemberships,
   setGroupsEnabled, resetDefaultTablesDB, setModelsMode, setViewAllCatalogs, setEnableHighlightHelpers, setUseMemory, addMemory, setCustomCSS, setEnableStyleCustomization, setEnableUserDebugTools, setEnableReviews, setMetadataHash, setMetadataProcessingCache, clearMetadataProcessingCache,
-  updateManualContextSelection
+  updateManualContextSelection, setCurrentEmail
 } = settingsSlice.actions
 
 export default settingsSlice.reducer

web/src/state/store.ts

@@ -16,6 +16,7 @@ import cache from './cache/reducer'
 import notifications from './notifications/reducer'
 import { userStateApi } from '../app/api/userStateApi'
 import { get } from 'lodash'
+import { getParsedIframeInfo } from '../helpers/origin'
 
 const combinedReducer = combineReducers({
   chat,
@@ -502,13 +503,25 @@ const migrations = {
     newState.settings.enableReviews = true
     return newState
   },
+  50: (state: RootState) => {
+    let newState = {...state}
+    // Refresh metadataProcessingCache
+    newState.settings.currentEmail = newState.auth?.email
+    return newState
+  },
+}
+
+const BLACKLIST = ['billing', 'cache', userStateApi.reducerPath]
+const isEmbedded = getParsedIframeInfo().isEmbedded as unknown === 'true'
+if (isEmbedded) {
+  BLACKLIST.push('auth')
 }
 
 const persistConfig = {
   key: 'root',
-  version: 49,
+  version: 50,
   storage,
-  blacklist: ['billing', 'cache', userStateApi.reducerPath],
+  blacklist: BLACKLIST,
   // @ts-ignore
   migrate: createMigrate(migrations, { debug: true }),
 };