fix(ui) Fix oauth redirect on connector modal (#3926)

* Fix oauth redirect on connector modal

* Fix lint

---------

Co-authored-by: Theodore Li <theo@sim.ai>

Author: TheodoreSpeaks

apps/sim/app/workspace/[workspaceId]/components/oauth-modal.tsx

@@ -16,7 +16,7 @@ import {
 } from '@/components/emcn'
 import { client, useSession } from '@/lib/auth/auth-client'
 import type { OAuthReturnContext } from '@/lib/credentials/client-state'
-import { writeOAuthReturnContext } from '@/lib/credentials/client-state'
+import { ADD_CONNECTOR_SEARCH_PARAM, writeOAuthReturnContext } from '@/lib/credentials/client-state'
 import {
   getCanonicalScopesForProvider,
   getProviderIdFromServiceId,
@@ -59,8 +59,8 @@ type OAuthModalConnectProps = OAuthModalBaseProps & {
   workspaceId: string
   credentialCount: number
 } & (
-    | { workflowId: string; knowledgeBaseId?: never }
-    | { workflowId?: never; knowledgeBaseId: string }
+    | { workflowId: string; knowledgeBaseId?: never; connectorType?: never }
+    | { workflowId?: never; knowledgeBaseId: string; connectorType?: string }
   )
 
 interface OAuthModalReauthorizeProps extends OAuthModalBaseProps {
@@ -81,6 +81,7 @@ export function OAuthModal(props: OAuthModalProps) {
   const workspaceId = isConnect ? props.workspaceId : ''
   const workflowId = isConnect ? props.workflowId : undefined
   const knowledgeBaseId = isConnect ? props.knowledgeBaseId : undefined
+  const connectorType = isConnect ? props.connectorType : undefined
   const toolName = !isConnect ? props.toolName : ''
   const requiredScopes = !isConnect ? (props.requiredScopes ?? EMPTY_SCOPES) : EMPTY_SCOPES
   const newScopes = !isConnect ? (props.newScopes ?? EMPTY_SCOPES) : EMPTY_SCOPES
@@ -172,7 +173,7 @@ export function OAuthModal(props: OAuthModalProps) {
         }
 
         const returnContext: OAuthReturnContext = knowledgeBaseId
-          ? { ...baseContext, origin: 'kb-connectors' as const, knowledgeBaseId }
+          ? { ...baseContext, origin: 'kb-connectors' as const, knowledgeBaseId, connectorType }
           : { ...baseContext, origin: 'workflow' as const, workflowId: workflowId! }
 
         writeOAuthReturnContext(returnContext)
@@ -205,7 +206,11 @@ export function OAuthModal(props: OAuthModalProps) {
         return
       }
 
-      await client.oauth2.link({ providerId, callbackURL: window.location.href })
+      const callbackURL = new URL(window.location.href)
+      if (connectorType) {
+        callbackURL.searchParams.set(ADD_CONNECTOR_SEARCH_PARAM, connectorType)
+      }
+      await client.oauth2.link({ providerId, callbackURL: callbackURL.toString() })
       handleClose()
     } catch (err) {
       logger.error('Failed to initiate OAuth connection', { error: err })

apps/sim/app/workspace/[workspaceId]/knowledge/[id]/base.tsx

@@ -4,7 +4,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import { createLogger } from '@sim/logger'
 import { format } from 'date-fns'
 import { AlertCircle, Loader2, Pencil, Plus, Tag, X } from 'lucide-react'
-import { useParams, useRouter } from 'next/navigation'
+import { useParams, usePathname, useRouter, useSearchParams } from 'next/navigation'
 import { usePostHog } from 'posthog-js/react'
 import {
   Badge,
@@ -25,6 +25,7 @@ import {
 import { Database, DatabaseX } from '@/components/emcn/icons'
 import { SearchHighlight } from '@/components/ui/search-highlight'
 import { cn } from '@/lib/core/utils/cn'
+import { ADD_CONNECTOR_SEARCH_PARAM } from '@/lib/credentials/client-state'
 import { ALL_TAG_SLOTS, type AllTagSlot, getFieldTypeForSlot } from '@/lib/knowledge/constants'
 import type { DocumentSortField, SortOrder } from '@/lib/knowledge/documents/types'
 import { type FilterFieldType, getOperatorsForFieldType } from '@/lib/knowledge/filters/types'
@@ -192,6 +193,10 @@ export function KnowledgeBase({
 }: KnowledgeBaseProps) {
   const params = useParams()
   const workspaceId = propWorkspaceId || (params.workspaceId as string)
+  const router = useRouter()
+  const searchParams = useSearchParams()
+  const pathname = usePathname()
+  const addConnectorParam = searchParams.get(ADD_CONNECTOR_SEARCH_PARAM)
   const posthog = usePostHog()
 
   useEffect(() => {
@@ -278,7 +283,29 @@ export function KnowledgeBase({
   const [contextMenuDocument, setContextMenuDocument] = useState<DocumentData | null>(null)
   const [showRenameModal, setShowRenameModal] = useState(false)
   const [documentToRename, setDocumentToRename] = useState<DocumentData | null>(null)
-  const [showAddConnectorModal, setShowAddConnectorModal] = useState(false)
+  const showAddConnectorModal = addConnectorParam != null
+  const searchParamsRef = useRef(searchParams)
+  searchParamsRef.current = searchParams
+  const updateAddConnectorParam = useCallback(
+    (value: string | null) => {
+      const current = searchParamsRef.current
+      const currentValue = current.get(ADD_CONNECTOR_SEARCH_PARAM)
+      if (value === currentValue || (value === null && currentValue === null)) return
+      const next = new URLSearchParams(current.toString())
+      if (value === null) {
+        next.delete(ADD_CONNECTOR_SEARCH_PARAM)
+      } else {
+        next.set(ADD_CONNECTOR_SEARCH_PARAM, value)
+      }
+      const qs = next.toString()
+      router.replace(qs ? `${pathname}?${qs}` : pathname, { scroll: false })
+    },
+    [pathname, router]
+  )
+  const setShowAddConnectorModal = useCallback(
+    (open: boolean) => updateAddConnectorParam(open ? '' : null),
+    [updateAddConnectorParam]
+  )
 
   const {
     isOpen: isContextMenuOpen,
@@ -340,8 +367,6 @@ export function KnowledgeBase({
     prevHadSyncingRef.current = hasSyncingConnectors
   }, [hasSyncingConnectors, refreshKnowledgeBase, refreshDocuments])
 
-  const router = useRouter()
-
   const knowledgeBaseName = knowledgeBase?.name || passedKnowledgeBaseName || 'Knowledge Base'
   const error = knowledgeBaseError || documentsError
 
@@ -1254,7 +1279,13 @@ export function KnowledgeBase({
       />
 
       {showAddConnectorModal && (
-        <AddConnectorModal open onOpenChange={setShowAddConnectorModal} knowledgeBaseId={id} />
+        <AddConnectorModal
+          open
+          onOpenChange={setShowAddConnectorModal}
+          onConnectorTypeChange={updateAddConnectorParam}
+          knowledgeBaseId={id}
+          initialConnectorType={addConnectorParam || undefined}
+        />
       )}
 
       {documentToRename && (

apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/add-connector-modal/add-connector-modal.tsx

@@ -44,14 +44,22 @@ const CONNECTOR_ENTRIES = Object.entries(CONNECTOR_REGISTRY)
 interface AddConnectorModalProps {
   open: boolean
   onOpenChange: (open: boolean) => void
+  onConnectorTypeChange?: (connectorType: string | null) => void
   knowledgeBaseId: string
+  initialConnectorType?: string | null
 }
 
 type Step = 'select-type' | 'configure'
 
-export function AddConnectorModal({ open, onOpenChange, knowledgeBaseId }: AddConnectorModalProps) {
-  const [step, setStep] = useState<Step>('select-type')
-  const [selectedType, setSelectedType] = useState<string | null>(null)
+export function AddConnectorModal({
+  open,
+  onOpenChange,
+  onConnectorTypeChange,
+  knowledgeBaseId,
+  initialConnectorType,
+}: AddConnectorModalProps) {
+  const [step, setStep] = useState<Step>(() => (initialConnectorType ? 'configure' : 'select-type'))
+  const [selectedType, setSelectedType] = useState<string | null>(initialConnectorType ?? null)
   const [sourceConfig, setSourceConfig] = useState<Record<string, string>>({})
   const [syncInterval, setSyncInterval] = useState(1440)
   const [selectedCredentialId, setSelectedCredentialId] = useState<string | null>(null)
@@ -151,6 +159,7 @@ export function AddConnectorModal({ open, onOpenChange, knowledgeBaseId }: AddCo
     setError(null)
     setSearchTerm('')
     setStep('configure')
+    onConnectorTypeChange?.(type)
   }
 
   const handleFieldChange = useCallback(
@@ -286,7 +295,10 @@ export function AddConnectorModal({ open, onOpenChange, knowledgeBaseId }: AddCo
               <Button
                 variant='ghost'
                 className='mr-2 h-6 w-6 p-0'
-                onClick={() => setStep('select-type')}
+                onClick={() => {
+                  setStep('select-type')
+                  onConnectorTypeChange?.('')
+                }}
               >
                 <ArrowLeft className='h-4 w-4' />
               </Button>
@@ -565,6 +577,7 @@ export function AddConnectorModal({ open, onOpenChange, knowledgeBaseId }: AddCo
             workspaceId={workspaceId}
             knowledgeBaseId={knowledgeBaseId}
             credentialCount={credentials.length}
+            connectorType={selectedType ?? undefined}
           />
         )}
     </>

apps/sim/hooks/use-oauth-return.ts

@@ -4,6 +4,7 @@ import { useEffect, useRef } from 'react'
 import { useParams, useRouter } from 'next/navigation'
 import { toast } from '@/components/emcn'
 import {
+  ADD_CONNECTOR_SEARCH_PARAM,
   consumeOAuthReturnContext,
   type OAuthReturnContext,
   readOAuthReturnContext,
@@ -98,7 +99,11 @@ export function useOAuthReturnRouter() {
       try {
         sessionStorage.removeItem(SETTINGS_RETURN_URL_KEY)
       } catch {}
-      router.replace(`/workspace/${workspaceId}/knowledge/${ctx.knowledgeBaseId}`)
+      const kbUrl = `/workspace/${workspaceId}/knowledge/${ctx.knowledgeBaseId}`
+      const connectorParam = ctx.connectorType
+        ? `?${ADD_CONNECTOR_SEARCH_PARAM}=${encodeURIComponent(ctx.connectorType)}`
+        : ''
+      router.replace(`${kbUrl}${connectorParam}`)
       return
     }
   }, [router, workspaceId])

apps/sim/lib/credentials/client-state.ts

@@ -91,6 +91,8 @@ export function clearPendingCredentialCreateRequest() {
   window.sessionStorage.removeItem(PENDING_CREDENTIAL_CREATE_REQUEST_KEY)
 }
 
+export const ADD_CONNECTOR_SEARCH_PARAM = 'addConnector' as const
+
 const OAUTH_RETURN_CONTEXT_KEY = 'sim.oauth-return-context'
 
 export type OAuthReturnOrigin = 'workflow' | 'integrations' | 'kb-connectors'
@@ -116,6 +118,7 @@ interface OAuthReturnIntegrations extends OAuthReturnBase {
 interface OAuthReturnKBConnectors extends OAuthReturnBase {
   origin: 'kb-connectors'
   knowledgeBaseId: string
+  connectorType?: string
 }
 
 export type OAuthReturnContext =