From 8906439a7e0eb2df53501e3505d97b9aec6d2b41 Mon Sep 17 00:00:00 2001
From: Vikhyath Mondreti <vikhyathvikku@gmail.com>
Date: Sun, 15 Mar 2026 01:01:11 -0700
Subject: [PATCH] fix(notifications): credentials connection notifs showing up
 in right resource (#3599)

* fix(notifications): credentials connection notifs showing up in right resource

* fix new label

* address comments

* reset ref correctly:
---
 .../[workspaceId]/knowledge/[id]/base.tsx     |   2 +
 .../add-connector-modal.tsx                   |  27 +++-
 .../connectors-section/connectors-section.tsx |  19 ++-
 .../app/workspace/[workspaceId]/layout.tsx    |   5 +-
 .../integrations/integrations-manager.tsx     |  73 +++++++--
 .../credential-selector.tsx                   |   5 +-
 .../[workspaceId]/w/[workflowId]/workflow.tsx |  64 +-------
 apps/sim/hooks/use-oauth-return.ts            | 147 ++++++++++++++++++
 apps/sim/lib/credentials/client-state.ts      |  59 +++++++
 9 files changed, 315 insertions(+), 86 deletions(-)
 create mode 100644 apps/sim/hooks/use-oauth-return.ts

diff --git a/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/base.tsx b/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/base.tsx
index 369b90bb2a..b813178612 100644
--- a/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/base.tsx
+++ b/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/base.tsx
@@ -72,6 +72,7 @@ import {
   useUpdateKnowledgeBase,
 } from '@/hooks/queries/kb/knowledge'
 import { useInlineRename } from '@/hooks/use-inline-rename'
+import { useOAuthReturnForKBConnectors } from '@/hooks/use-oauth-return'
 
 const logger = createLogger('KnowledgeBase')
 
@@ -189,6 +190,7 @@ export function KnowledgeBase({
 }: KnowledgeBaseProps) {
   const params = useParams()
   const workspaceId = propWorkspaceId || (params.workspaceId as string)
+  useOAuthReturnForKBConnectors(id)
   const { removeKnowledgeBase } = useKnowledgeBasesList(workspaceId, { enabled: false })
   const userPermissions = useUserPermissionsContext()
 
diff --git a/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/add-connector-modal/add-connector-modal.tsx b/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/add-connector-modal/add-connector-modal.tsx
index 5d3a7a9337..17048ab27e 100644
--- a/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/add-connector-modal/add-connector-modal.tsx
+++ b/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/add-connector-modal/add-connector-modal.tsx
@@ -20,6 +20,7 @@ import {
   Tooltip,
 } from '@/components/emcn'
 import { useSession } from '@/lib/auth/auth-client'
+import { consumeOAuthReturnContext, writeOAuthReturnContext } from '@/lib/credentials/client-state'
 import {
   getCanonicalScopesForProvider,
   getProviderIdFromServiceId,
@@ -288,8 +289,25 @@ export function AddConnectorModal({ open, onOpenChange, knowledgeBaseId }: AddCo
       return
     }
 
+    writeOAuthReturnContext({
+      origin: 'kb-connectors',
+      knowledgeBaseId,
+      displayName,
+      providerId: connectorProviderId,
+      preCount: credentials.length,
+      workspaceId,
+      requestedAt: Date.now(),
+    })
+
     setShowOAuthModal(true)
-  }, [connectorConfig, connectorProviderId, workspaceId, session?.user?.name])
+  }, [
+    connectorConfig,
+    connectorProviderId,
+    workspaceId,
+    session?.user?.name,
+    knowledgeBaseId,
+    credentials.length,
+  ])
 
   const filteredEntries = useMemo(() => {
     const term = searchTerm.toLowerCase().trim()
@@ -575,11 +593,14 @@ export function AddConnectorModal({ open, onOpenChange, knowledgeBaseId }: AddCo
       {connectorConfig && connectorConfig.auth.mode === 'oauth' && connectorProviderId && (
         <OAuthRequiredModal
           isOpen={showOAuthModal}
-          onClose={() => setShowOAuthModal(false)}
+          onClose={() => {
+            consumeOAuthReturnContext()
+            setShowOAuthModal(false)
+          }}
           provider={connectorProviderId}
           toolName={connectorConfig.name}
           requiredScopes={getCanonicalScopesForProvider(connectorProviderId)}
-          newScopes={connectorConfig.auth.requiredScopes || []}
+          newScopes={[]}
           serviceId={connectorConfig.auth.provider}
         />
       )}
diff --git a/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/connectors-section/connectors-section.tsx b/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/connectors-section/connectors-section.tsx
index 58884a0de3..3abc3d5082 100644
--- a/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/connectors-section/connectors-section.tsx
+++ b/apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/connectors-section/connectors-section.tsx
@@ -27,6 +27,7 @@ import {
   Tooltip,
 } from '@/components/emcn'
 import { cn } from '@/lib/core/utils/cn'
+import { consumeOAuthReturnContext, writeOAuthReturnContext } from '@/lib/credentials/client-state'
 import {
   getCanonicalScopesForProvider,
   getProviderIdFromServiceId,
@@ -444,7 +445,18 @@ function ConnectorCard({
             {canEdit && (
               <Button
                 variant='active'
-                onClick={() => setShowOAuthModal(true)}
+                onClick={() => {
+                  writeOAuthReturnContext({
+                    origin: 'kb-connectors',
+                    knowledgeBaseId,
+                    displayName: connectorDef?.name ?? connector.connectorType,
+                    providerId: providerId!,
+                    preCount: credentials?.length ?? 0,
+                    workspaceId,
+                    requestedAt: Date.now(),
+                  })
+                  setShowOAuthModal(true)
+                }}
                 className='w-full px-[8px] py-[4px] font-medium text-[12px]'
               >
                 Update access
@@ -463,7 +475,10 @@ function ConnectorCard({
       {showOAuthModal && serviceId && providerId && (
         <OAuthRequiredModal
           isOpen={showOAuthModal}
-          onClose={() => setShowOAuthModal(false)}
+          onClose={() => {
+            consumeOAuthReturnContext()
+            setShowOAuthModal(false)
+          }}
           provider={providerId as OAuthProvider}
           toolName={connectorDef?.name ?? connector.connectorType}
           requiredScopes={getCanonicalScopesForProvider(providerId)}
diff --git a/apps/sim/app/workspace/[workspaceId]/layout.tsx b/apps/sim/app/workspace/[workspaceId]/layout.tsx
index de24323dd8..5abbde90ab 100644
--- a/apps/sim/app/workspace/[workspaceId]/layout.tsx
+++ b/apps/sim/app/workspace/[workspaceId]/layout.tsx
@@ -1,5 +1,6 @@
 'use client'
 
+import { ToastProvider } from '@/components/emcn'
 import { GlobalCommandsProvider } from '@/app/workspace/[workspaceId]/providers/global-commands-provider'
 import { ProviderModelsLoader } from '@/app/workspace/[workspaceId]/providers/provider-models-loader'
 import { SettingsLoader } from '@/app/workspace/[workspaceId]/providers/settings-loader'
@@ -8,7 +9,7 @@ import { Sidebar } from '@/app/workspace/[workspaceId]/w/components/sidebar/side
 
 export default function WorkspaceLayout({ children }: { children: React.ReactNode }) {
   return (
-    <>
+    <ToastProvider>
       <SettingsLoader />
       <ProviderModelsLoader />
       <GlobalCommandsProvider>
@@ -25,6 +26,6 @@ export default function WorkspaceLayout({ children }: { children: React.ReactNod
           </WorkspacePermissionsProvider>
         </div>
       </GlobalCommandsProvider>
-    </>
+    </ToastProvider>
   )
 }
diff --git a/apps/sim/app/workspace/[workspaceId]/settings/components/integrations/integrations-manager.tsx b/apps/sim/app/workspace/[workspaceId]/settings/components/integrations/integrations-manager.tsx
index 32a95b03dd..7518dc6fed 100644
--- a/apps/sim/app/workspace/[workspaceId]/settings/components/integrations/integrations-manager.tsx
+++ b/apps/sim/app/workspace/[workspaceId]/settings/components/integrations/integrations-manager.tsx
@@ -1,6 +1,6 @@
 'use client'
 
-import { createElement, useCallback, useEffect, useMemo, useState } from 'react'
+import { createElement, useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import { createLogger } from '@sim/logger'
 import { AlertTriangle, Check, Clipboard, Plus, Search, Share2 } from 'lucide-react'
 import { useParams } from 'next/navigation'
@@ -28,6 +28,7 @@ import {
   PENDING_CREDENTIAL_CREATE_REQUEST_EVENT,
   type PendingCredentialCreateRequest,
   readPendingCredentialCreateRequest,
+  writeOAuthReturnContext,
 } from '@/lib/credentials/client-state'
 import {
   getCanonicalScopesForProvider,
@@ -54,6 +55,7 @@ import {
   useOAuthConnections,
 } from '@/hooks/queries/oauth/oauth-connections'
 import { useWorkspacePermissionsQuery } from '@/hooks/queries/workspace'
+import { useOAuthReturnRouter } from '@/hooks/use-oauth-return'
 
 const logger = createLogger('IntegrationsManager')
 
@@ -66,6 +68,8 @@ export function IntegrationsManager() {
   const params = useParams()
   const workspaceId = (params?.workspaceId as string) || ''
 
+  useOAuthReturnRouter()
+
   const [searchTerm, setSearchTerm] = useState('')
   const [selectedCredentialId, setSelectedCredentialId] = useState<string | null>(null)
   const [memberRole, setMemberRole] = useState<WorkspaceCredentialRole>('admin')
@@ -84,6 +88,11 @@ export function IntegrationsManager() {
   const [showDeleteConfirmDialog, setShowDeleteConfirmDialog] = useState(false)
   const [deleteError, setDeleteError] = useState<string | null>(null)
   const [showUnsavedChangesAlert, setShowUnsavedChangesAlert] = useState(false)
+  const pendingReturnOriginRef = useRef<
+    | { type: 'workflow'; workflowId: string }
+    | { type: 'kb-connectors'; knowledgeBaseId: string }
+    | undefined
+  >(undefined)
   const { data: session } = useSession()
   const currentUserId = session?.user?.id || ''
 
@@ -278,6 +287,8 @@ export function IntegrationsManager() {
 
       if (request.type !== 'oauth') return
 
+      pendingReturnOriginRef.current = request.returnOrigin
+
       setShowCreateModal(true)
       setShowCreateOAuthRequiredModal(false)
       setCreateError(null)
@@ -350,6 +361,7 @@ export function IntegrationsManager() {
     setCreateOAuthProviderId('')
     setCreateError(null)
     setShowCreateOAuthRequiredModal(false)
+    pendingReturnOriginRef.current = undefined
   }
 
   const handleSelectCredential = (credential: WorkspaceCredential) => {
@@ -397,15 +409,42 @@ export function IntegrationsManager() {
         }),
       })
 
-      window.sessionStorage.setItem(
-        'sim.oauth-connect-pending',
-        JSON.stringify({
+      const oauthPreCount = credentials.filter(
+        (c) => c.type === 'oauth' && c.providerId === selectedOAuthService.providerId
+      ).length
+      const returnOrigin = pendingReturnOriginRef.current
+      pendingReturnOriginRef.current = undefined
+
+      if (returnOrigin?.type === 'workflow') {
+        writeOAuthReturnContext({
+          origin: 'workflow',
+          workflowId: returnOrigin.workflowId,
           displayName,
           providerId: selectedOAuthService.providerId,
-          preCount: credentials.filter((c) => c.type === 'oauth').length,
+          preCount: oauthPreCount,
           workspaceId,
+          requestedAt: Date.now(),
         })
-      )
+      } else if (returnOrigin?.type === 'kb-connectors') {
+        writeOAuthReturnContext({
+          origin: 'kb-connectors',
+          knowledgeBaseId: returnOrigin.knowledgeBaseId,
+          displayName,
+          providerId: selectedOAuthService.providerId,
+          preCount: oauthPreCount,
+          workspaceId,
+          requestedAt: Date.now(),
+        })
+      } else {
+        writeOAuthReturnContext({
+          origin: 'integrations',
+          displayName,
+          providerId: selectedOAuthService.providerId,
+          preCount: oauthPreCount,
+          workspaceId,
+          requestedAt: Date.now(),
+        })
+      }
 
       await connectOAuthService.mutateAsync({
         providerId: selectedOAuthService.providerId,
@@ -512,16 +551,18 @@ export function IntegrationsManager() {
         }),
       })
 
-      window.sessionStorage.setItem(
-        'sim.oauth-connect-pending',
-        JSON.stringify({
-          displayName: selectedCredential.displayName,
-          providerId: selectedCredential.providerId,
-          preCount: credentials.filter((c) => c.type === 'oauth').length,
-          workspaceId,
-          reconnect: true,
-        })
-      )
+      const oauthPreCount = credentials.filter(
+        (c) => c.type === 'oauth' && c.providerId === selectedCredential.providerId
+      ).length
+      writeOAuthReturnContext({
+        origin: 'integrations',
+        displayName: selectedCredential.displayName,
+        providerId: selectedCredential.providerId,
+        preCount: oauthPreCount,
+        workspaceId,
+        reconnect: true,
+        requestedAt: Date.now(),
+      })
 
       await connectOAuthService.mutateAsync({
         providerId: selectedCredential.providerId,
diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/credential-selector/credential-selector.tsx b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/credential-selector/credential-selector.tsx
index 8d5d52637a..43b0ee3631 100644
--- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/credential-selector/credential-selector.tsx
+++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/credential-selector/credential-selector.tsx
@@ -207,10 +207,13 @@ export function CredentialSelector({
       serviceId,
       requiredScopes: getCanonicalScopesForProvider(effectiveProviderId),
       requestedAt: Date.now(),
+      returnOrigin: activeWorkflowId
+        ? { type: 'workflow', workflowId: activeWorkflowId }
+        : undefined,
     })
 
     navigateToSettings({ section: 'integrations' })
-  }, [workspaceId, effectiveProviderId, serviceId])
+  }, [workspaceId, effectiveProviderId, serviceId, activeWorkflowId])
 
   const getProviderIcon = useCallback((providerName: OAuthProvider) => {
     const { baseProvider } = parseProvider(providerName)
diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/workflow.tsx b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/workflow.tsx
index f026270c89..c1f78d9662 100644
--- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/workflow.tsx
+++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/workflow.tsx
@@ -73,6 +73,7 @@ import { useWorkspaceEnvironment } from '@/hooks/queries/environment'
 import { useAutoConnect, useSnapToGridSize } from '@/hooks/queries/general-settings'
 import { useCanvasViewport } from '@/hooks/use-canvas-viewport'
 import { useCollaborativeWorkflow } from '@/hooks/use-collaborative-workflow'
+import { useOAuthReturnForWorkflow } from '@/hooks/use-oauth-return'
 import { useStreamCleanup } from '@/hooks/use-stream-cleanup'
 import { useCanvasModeStore } from '@/stores/canvas-mode'
 import { useChatStore } from '@/stores/chat/store'
@@ -268,68 +269,7 @@ const WorkflowContent = React.memo(
 
     const addNotification = useNotificationStore((state) => state.addNotification)
 
-    useEffect(() => {
-      const OAUTH_CONNECT_PENDING_KEY = 'sim.oauth-connect-pending'
-      const pending = window.sessionStorage.getItem(OAUTH_CONNECT_PENDING_KEY)
-      if (!pending) return
-      window.sessionStorage.removeItem(OAUTH_CONNECT_PENDING_KEY)
-
-      ;(async () => {
-        try {
-          const {
-            displayName,
-            providerId,
-            preCount,
-            workspaceId: wsId,
-            reconnect,
-          } = JSON.parse(pending) as {
-            displayName: string
-            providerId: string
-            preCount: number
-            workspaceId: string
-            reconnect?: boolean
-          }
-
-          if (reconnect) {
-            addNotification({
-              level: 'info',
-              message: `"${displayName}" reconnected successfully.`,
-            })
-            window.dispatchEvent(
-              new CustomEvent('oauth-credentials-updated', {
-                detail: { providerId, workspaceId: wsId },
-              })
-            )
-            return
-          }
-
-          const response = await fetch(
-            `/api/credentials?workspaceId=${encodeURIComponent(wsId)}&type=oauth`
-          )
-          const data = response.ok ? await response.json() : { credentials: [] }
-          const oauthCredentials = (data.credentials ?? []) as Array<{
-            displayName: string
-            providerId: string | null
-          }>
-
-          if (oauthCredentials.length > preCount) {
-            addNotification({
-              level: 'info',
-              message: `"${displayName}" credential connected successfully.`,
-            })
-          } else {
-            const existing = oauthCredentials.find((c) => c.providerId === providerId)
-            const existingName = existing?.displayName || displayName
-            addNotification({
-              level: 'info',
-              message: `This account is already connected as "${existingName}".`,
-            })
-          }
-        } catch {
-          // Ignore malformed sessionStorage data
-        }
-      })()
-    }, [])
+    useOAuthReturnForWorkflow(workflowIdParam)
 
     const {
       workflows,
diff --git a/apps/sim/hooks/use-oauth-return.ts b/apps/sim/hooks/use-oauth-return.ts
new file mode 100644
index 0000000000..ad64918b87
--- /dev/null
+++ b/apps/sim/hooks/use-oauth-return.ts
@@ -0,0 +1,147 @@
+'use client'
+
+import { useEffect, useRef } from 'react'
+import { useParams, useRouter } from 'next/navigation'
+import { toast } from '@/components/emcn'
+import {
+  consumeOAuthReturnContext,
+  type OAuthReturnContext,
+  readOAuthReturnContext,
+} from '@/lib/credentials/client-state'
+import { useNotificationStore } from '@/stores/notifications/store'
+
+const OAUTH_CREDENTIAL_UPDATED_EVENT = 'oauth-credentials-updated'
+const SETTINGS_RETURN_URL_KEY = 'settings-return-url'
+const CONTEXT_MAX_AGE_MS = 15 * 60 * 1000
+
+async function resolveOAuthMessage(ctx: OAuthReturnContext): Promise<string> {
+  if (ctx.reconnect) {
+    return `"${ctx.displayName}" reconnected successfully.`
+  }
+
+  try {
+    const response = await fetch(
+      `/api/credentials?workspaceId=${encodeURIComponent(ctx.workspaceId)}&type=oauth`
+    )
+    const data = response.ok ? await response.json() : { credentials: [] }
+    const oauthCredentials = (data.credentials ?? []) as Array<{
+      displayName: string
+      providerId: string | null
+    }>
+
+    const forProvider = oauthCredentials.filter((c) => c.providerId === ctx.providerId)
+    if (forProvider.length > ctx.preCount) {
+      return `"${ctx.displayName}" credential connected successfully.`
+    }
+
+    const existing = forProvider[0]
+    return `This account is already connected as "${existing?.displayName || ctx.displayName}".`
+  } catch {
+    return `"${ctx.displayName}" credential connected successfully.`
+  }
+}
+
+function dispatchCredentialUpdate(ctx: OAuthReturnContext) {
+  window.dispatchEvent(
+    new CustomEvent(OAUTH_CREDENTIAL_UPDATED_EVENT, {
+      detail: { providerId: ctx.providerId, workspaceId: ctx.workspaceId },
+    })
+  )
+}
+
+/**
+ * Post-OAuth router for the integrations page.
+ *
+ * After OAuth, Better Auth redirects back to `callbackURL` which is the integrations page.
+ * This hook reads the stored return context to determine the original initiator:
+ *
+ * - `integrations`: Stay on this page, show a toast notification.
+ * - `workflow`: Redirect to the specific workflow. The workflow page picks up the context.
+ * - `kb-connectors`: Redirect to the KB page. The KB page picks up the context.
+ */
+export function useOAuthReturnRouter() {
+  const router = useRouter()
+  const params = useParams()
+  const workspaceId = params.workspaceId as string
+  const handledRef = useRef(false)
+
+  useEffect(() => {
+    if (handledRef.current) return
+    const ctx = readOAuthReturnContext()
+    if (!ctx) return
+    if (Date.now() - ctx.requestedAt > CONTEXT_MAX_AGE_MS) {
+      consumeOAuthReturnContext()
+      return
+    }
+
+    handledRef.current = true
+
+    if (ctx.origin === 'integrations') {
+      consumeOAuthReturnContext()
+      void (async () => {
+        const message = await resolveOAuthMessage(ctx)
+        toast.success(message, { duration: 5000 })
+        dispatchCredentialUpdate(ctx)
+      })()
+      return
+    }
+
+    if (ctx.origin === 'workflow') {
+      try {
+        sessionStorage.removeItem(SETTINGS_RETURN_URL_KEY)
+      } catch {}
+      router.replace(`/workspace/${workspaceId}/w/${ctx.workflowId}`)
+      return
+    }
+
+    if (ctx.origin === 'kb-connectors') {
+      try {
+        sessionStorage.removeItem(SETTINGS_RETURN_URL_KEY)
+      } catch {}
+      router.replace(`/workspace/${workspaceId}/knowledge/${ctx.knowledgeBaseId}`)
+      return
+    }
+  }, [router, workspaceId])
+}
+
+/**
+ * Post-OAuth handler for workflow pages.
+ * Consumes the return context and shows a workflow-scoped notification.
+ */
+export function useOAuthReturnForWorkflow(workflowId: string) {
+  const addNotification = useNotificationStore((state) => state.addNotification)
+
+  useEffect(() => {
+    const ctx = readOAuthReturnContext()
+    if (!ctx || ctx.origin !== 'workflow') return
+    if (ctx.workflowId !== workflowId) return
+    consumeOAuthReturnContext()
+    if (Date.now() - ctx.requestedAt > CONTEXT_MAX_AGE_MS) return
+
+    void (async () => {
+      const message = await resolveOAuthMessage(ctx)
+      addNotification({ level: 'info', message, workflowId })
+      dispatchCredentialUpdate(ctx)
+    })()
+  }, [workflowId, addNotification])
+}
+
+/**
+ * Post-OAuth handler for KB connectors pages.
+ * Consumes the return context and shows a toast notification.
+ */
+export function useOAuthReturnForKBConnectors(knowledgeBaseId: string) {
+  useEffect(() => {
+    const ctx = readOAuthReturnContext()
+    if (!ctx || ctx.origin !== 'kb-connectors') return
+    if (ctx.knowledgeBaseId !== knowledgeBaseId) return
+    consumeOAuthReturnContext()
+    if (Date.now() - ctx.requestedAt > CONTEXT_MAX_AGE_MS) return
+
+    void (async () => {
+      const message = await resolveOAuthMessage(ctx)
+      toast.success(message, { duration: 5000 })
+      dispatchCredentialUpdate(ctx)
+    })()
+  }, [knowledgeBaseId])
+}
diff --git a/apps/sim/lib/credentials/client-state.ts b/apps/sim/lib/credentials/client-state.ts
index 9a333e5ea7..a7216677c4 100644
--- a/apps/sim/lib/credentials/client-state.ts
+++ b/apps/sim/lib/credentials/client-state.ts
@@ -21,6 +21,15 @@ interface PendingOAuthCredentialCreateRequest {
   serviceId: string
   requiredScopes: string[]
   requestedAt: number
+  returnOrigin?:
+    | {
+        type: 'workflow'
+        workflowId: string
+      }
+    | {
+        type: 'kb-connectors'
+        knowledgeBaseId: string
+      }
 }
 
 interface PendingSecretCredentialCreateRequest {
@@ -81,3 +90,53 @@ export function clearPendingCredentialCreateRequest() {
   if (typeof window === 'undefined') return
   window.sessionStorage.removeItem(PENDING_CREDENTIAL_CREATE_REQUEST_KEY)
 }
+
+const OAUTH_RETURN_CONTEXT_KEY = 'sim.oauth-return-context'
+
+export type OAuthReturnOrigin = 'workflow' | 'integrations' | 'kb-connectors'
+
+interface OAuthReturnBase {
+  displayName: string
+  providerId: string
+  preCount: number
+  workspaceId: string
+  reconnect?: boolean
+  requestedAt: number
+}
+
+interface OAuthReturnWorkflow extends OAuthReturnBase {
+  origin: 'workflow'
+  workflowId: string
+}
+
+interface OAuthReturnIntegrations extends OAuthReturnBase {
+  origin: 'integrations'
+}
+
+interface OAuthReturnKBConnectors extends OAuthReturnBase {
+  origin: 'kb-connectors'
+  knowledgeBaseId: string
+}
+
+export type OAuthReturnContext =
+  | OAuthReturnWorkflow
+  | OAuthReturnIntegrations
+  | OAuthReturnKBConnectors
+
+export function writeOAuthReturnContext(ctx: OAuthReturnContext) {
+  if (typeof window === 'undefined') return
+  window.sessionStorage.setItem(OAUTH_RETURN_CONTEXT_KEY, JSON.stringify(ctx))
+}
+
+export function readOAuthReturnContext(): OAuthReturnContext | null {
+  if (typeof window === 'undefined') return null
+  return parseJson<OAuthReturnContext>(window.sessionStorage.getItem(OAUTH_RETURN_CONTEXT_KEY))
+}
+
+export function consumeOAuthReturnContext(): OAuthReturnContext | null {
+  const ctx = readOAuthReturnContext()
+  if (ctx) {
+    window.sessionStorage.removeItem(OAUTH_RETURN_CONTEXT_KEY)
+  }
+  return ctx
+}