quick nit

* add tenants on standby feature

* k

* fix alembic

* k

* k

backend/alembic/versions/4d58345da04a_lowercase_user_emails.py

@@ -5,7 +5,10 @@ Revises: f1ca58b2f2ec
 Create Date: 2025-01-29 07:48:46.784041
 
 """
+import logging
+from typing import cast
 from alembic import op
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.sql import text
 
 
@@ -15,21 +18,45 @@ down_revision = "f1ca58b2f2ec"
 branch_labels = None
 depends_on = None
 
+logger = logging.getLogger("alembic.runtime.migration")
+
 
 def upgrade() -> None:
-    # Get database connection
+    """Conflicts on lowercasing will result in the uppercased email getting a
+    unique integer suffix when converted to lowercase."""
+
     connection = op.get_bind()
 
-    # Update all user emails to lowercase
-    connection.execute(
-        text(
-            """
-            UPDATE "user"
-            SET email = LOWER(email)
-            WHERE email != LOWER(email)
-            """
-        )
-    )
+    # Fetch all user emails that are not already lowercase
+    user_emails = connection.execute(
+        text('SELECT id, email FROM "user" WHERE email != LOWER(email)')
+    ).fetchall()
+
+    for user_id, email in user_emails:
+        email = cast(str, email)
+        username, domain = email.rsplit("@", 1)
+        new_email = f"{username.lower()}@{domain.lower()}"
+        attempt = 1
+
+        while True:
+            try:
+                # Try updating the email
+                connection.execute(
+                    text('UPDATE "user" SET email = :new_email WHERE id = :user_id'),
+                    {"new_email": new_email, "user_id": user_id},
+                )
+                break  # Success, exit loop
+            except IntegrityError:
+                next_email = f"{username.lower()}_{attempt}@{domain.lower()}"
+                # Email conflict occurred, append `_1`, `_2`, etc., to the username
+                logger.warning(
+                    f"Conflict while lowercasing email: "
+                    f"old_email={email} "
+                    f"conflicting_email={new_email} "
+                    f"next_email={next_email}"
+                )
+                new_email = next_email
+                attempt += 1
 
 
 def downgrade() -> None:

backend/alembic_tenants/versions/3b45e0018bf1_add_new_available_tenant_table.py

@@ -0,0 +1,33 @@
+"""add new available tenant table
+
+Revision ID: 3b45e0018bf1
+Revises: ac842f85f932
+Create Date: 2025-03-06 09:55:18.229910
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "3b45e0018bf1"
+down_revision = "ac842f85f932"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Create new_available_tenant table
+    op.create_table(
+        "available_tenant",
+        sa.Column("tenant_id", sa.String(), nullable=False),
+        sa.Column("alembic_version", sa.String(), nullable=False),
+        sa.Column("date_created", sa.DateTime(), nullable=False),
+        sa.PrimaryKeyConstraint("tenant_id"),
+    )
+
+
+def downgrade() -> None:
+    # Drop new_available_tenant table
+    op.drop_table("available_tenant")

backend/ee/onyx/server/tenants/provisioning.py

@@ -28,11 +28,12 @@ from onyx.auth.users import exceptions
 from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL
 from onyx.configs.app_configs import DEV_MODE
 from onyx.configs.constants import MilestoneRecordType
+from onyx.db.engine import get_session_with_shared_schema
 from onyx.db.engine import get_session_with_tenant
-from onyx.db.engine import get_sqlalchemy_engine
 from onyx.db.llm import update_default_provider
 from onyx.db.llm import upsert_cloud_embedding_provider
 from onyx.db.llm import upsert_llm_provider
+from onyx.db.models import AvailableTenant
 from onyx.db.models import IndexModelStatus
 from onyx.db.models import SearchSettings
 from onyx.db.models import UserTenantMapping
@@ -62,42 +63,72 @@ async def get_or_provision_tenant(
     This function should only be called after we have verified we want this user's tenant to exist.
     It returns the tenant ID associated with the email, creating a new tenant if necessary.
     """
+    # Early return for non-multi-tenant mode
     if not MULTI_TENANT:
         return POSTGRES_DEFAULT_SCHEMA
 
     if referral_source and request:
         await submit_to_hubspot(email, referral_source, request)
 
+    # First, check if the user already has a tenant
+    tenant_id: str | None = None
     try:
         tenant_id = get_tenant_id_for_email(email)
+        return tenant_id
     except exceptions.UserNotExists:
-        # If tenant does not exist and in Multi tenant mode, provision a new tenant
-        try:
+        # User doesn't exist, so we need to create a new tenant or assign an existing one
+        pass
+
+    try:
+        # Try to get a pre-provisioned tenant
+        tenant_id = await get_available_tenant()
+
+        if tenant_id:
+            # If we have a pre-provisioned tenant, assign it to the user
+            await assign_tenant_to_user(tenant_id, email, referral_source)
+            logger.info(f"Assigned pre-provisioned tenant {tenant_id} to user {email}")
+            return tenant_id
+        else:
+            # If no pre-provisioned tenant is available, create a new one on-demand
             tenant_id = await create_tenant(email, referral_source)
-        except Exception as e:
-            logger.error(f"Tenant provisioning failed: {e}")
-            raise HTTPException(status_code=500, detail="Failed to provision tenant.")
+            return tenant_id
 
-    if not tenant_id:
+    except Exception as e:
+        # If we've encountered an error, log and raise an exception
+        error_msg = "Failed to provision tenant"
+        logger.error(error_msg, exc_info=e)
         raise HTTPException(
-            status_code=401, detail="User does not belong to an organization"
+            status_code=500,
+            detail="Failed to provision tenant. Please try again later.",
         )
 
-    return tenant_id
-
 
 async def create_tenant(email: str, referral_source: str | None = None) -> str:
+    """
+    Create a new tenant on-demand when no pre-provisioned tenants are available.
+    This is the fallback method when we can't use a pre-provisioned tenant.
+
+    """
     tenant_id = TENANT_ID_PREFIX + str(uuid.uuid4())
+    logger.info(f"Creating new tenant {tenant_id} for user {email}")
+
     try:
         # Provision tenant on data plane
         await provision_tenant(tenant_id, email)
-        # Notify control plane
-        if not DEV_MODE:
+
+        # Notify control plane if not already done in provision_tenant
+        if not DEV_MODE and referral_source:
             await notify_control_plane(tenant_id, email, referral_source)
+
     except Exception as e:
-        logger.error(f"Tenant provisioning failed: {e}")
-        await rollback_tenant_provisioning(tenant_id)
+        logger.exception(f"Tenant provisioning failed: {str(e)}")
+        # Attempt to rollback the tenant provisioning
+        try:
+            await rollback_tenant_provisioning(tenant_id)
+        except Exception:
+            logger.exception(f"Failed to rollback tenant provisioning for {tenant_id}")
         raise HTTPException(status_code=500, detail="Failed to provision tenant.")
+
     return tenant_id
 
 
@@ -111,54 +142,25 @@ async def provision_tenant(tenant_id: str, email: str) -> None:
         )
 
     logger.debug(f"Provisioning tenant {tenant_id} for user {email}")
-    token = None
 
     try:
+        # Create the schema for the tenant
         if not create_schema_if_not_exists(tenant_id):
             logger.debug(f"Created schema for tenant {tenant_id}")
         else:
             logger.debug(f"Schema already exists for tenant {tenant_id}")
 
-        token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+        # Set up the tenant with all necessary configurations
+        await setup_tenant(tenant_id)
 
-        # Await the Alembic migrations
-        await asyncio.to_thread(run_alembic_migrations, tenant_id)
-
-        with get_session_with_tenant(tenant_id=tenant_id) as db_session:
-            configure_default_api_keys(db_session)
-
-            current_search_settings = (
-                db_session.query(SearchSettings)
-                .filter_by(status=IndexModelStatus.FUTURE)
-                .first()
-            )
-            cohere_enabled = (
-                current_search_settings is not None
-                and current_search_settings.provider_type == EmbeddingProvider.COHERE
-            )
-            setup_onyx(db_session, tenant_id, cohere_enabled=cohere_enabled)
-
-        add_users_to_tenant([email], tenant_id)
-
-        with get_session_with_tenant(tenant_id=tenant_id) as db_session:
-            create_milestone_and_report(
-                user=None,
-                distinct_id=tenant_id,
-                event_type=MilestoneRecordType.TENANT_CREATED,
-                properties={
-                    "email": email,
-                },
-                db_session=db_session,
-            )
+        # Assign the tenant to the user
+        await assign_tenant_to_user(tenant_id, email)
 
     except Exception as e:
         logger.exception(f"Failed to create tenant {tenant_id}")
         raise HTTPException(
             status_code=500, detail=f"Failed to create tenant: {str(e)}"
         )
-    finally:
-        if token is not None:
-            CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
 
 
 async def notify_control_plane(
@@ -189,20 +191,74 @@ async def notify_control_plane(
 
 
 async def rollback_tenant_provisioning(tenant_id: str) -> None:
-    # Logic to rollback tenant provisioning on data plane
+    """
+    Logic to rollback tenant provisioning on data plane.
+    Handles each step independently to ensure maximum cleanup even if some steps fail.
+    """
     logger.info(f"Rolling back tenant provisioning for tenant_id: {tenant_id}")
-    try:
-        # Drop the tenant's schema to rollback provisioning
-        drop_schema(tenant_id)
 
-        # Remove tenant mapping
-        with Session(get_sqlalchemy_engine()) as db_session:
-            db_session.query(UserTenantMapping).filter(
-                UserTenantMapping.tenant_id == tenant_id
-            ).delete()
-            db_session.commit()
+    # Track if any part of the rollback fails
+    rollback_errors = []
+
+    # 1. Try to drop the tenant's schema
+    try:
+        drop_schema(tenant_id)
+        logger.info(f"Successfully dropped schema for tenant {tenant_id}")
     except Exception as e:
-        logger.error(f"Failed to rollback tenant provisioning: {e}")
+        error_msg = f"Failed to drop schema for tenant {tenant_id}: {str(e)}"
+        logger.error(error_msg)
+        rollback_errors.append(error_msg)
+
+    # 2. Try to remove tenant mapping
+    try:
+        with get_session_with_shared_schema() as db_session:
+            db_session.begin()
+            try:
+                db_session.query(UserTenantMapping).filter(
+                    UserTenantMapping.tenant_id == tenant_id
+                ).delete()
+                db_session.commit()
+                logger.info(
+                    f"Successfully removed user mappings for tenant {tenant_id}"
+                )
+            except Exception as e:
+                db_session.rollback()
+                raise e
+    except Exception as e:
+        error_msg = f"Failed to remove user mappings for tenant {tenant_id}: {str(e)}"
+        logger.error(error_msg)
+        rollback_errors.append(error_msg)
+
+    # 3. If this tenant was in the available tenants table, remove it
+    try:
+        with get_session_with_shared_schema() as db_session:
+            db_session.begin()
+            try:
+                available_tenant = (
+                    db_session.query(AvailableTenant)
+                    .filter(AvailableTenant.tenant_id == tenant_id)
+                    .first()
+                )
+
+                if available_tenant:
+                    db_session.delete(available_tenant)
+                    db_session.commit()
+                    logger.info(
+                        f"Removed tenant {tenant_id} from available tenants table"
+                    )
+            except Exception as e:
+                db_session.rollback()
+                raise e
+    except Exception as e:
+        error_msg = f"Failed to remove tenant {tenant_id} from available tenants table: {str(e)}"
+        logger.error(error_msg)
+        rollback_errors.append(error_msg)
+
+    # Log summary of rollback operation
+    if rollback_errors:
+        logger.error(f"Tenant rollback completed with {len(rollback_errors)} errors")
+    else:
+        logger.info(f"Tenant rollback completed successfully for tenant {tenant_id}")
 
 
 def configure_default_api_keys(db_session: Session) -> None:
@@ -399,3 +455,111 @@ def get_tenant_by_domain_from_control_plane(
     except Exception as e:
         logger.error(f"Error fetching tenant by domain: {str(e)}")
         return None
+
+
+async def get_available_tenant() -> str | None:
+    """
+    Get an available pre-provisioned tenant from the NewAvailableTenant table.
+    Returns the tenant_id if one is available, None otherwise.
+    Uses row-level locking to prevent race conditions when multiple processes
+    try to get an available tenant simultaneously.
+    """
+    if not MULTI_TENANT:
+        return None
+
+    with get_session_with_shared_schema() as db_session:
+        try:
+            db_session.begin()
+
+            # Get the oldest available tenant with FOR UPDATE lock to prevent race conditions
+            available_tenant = (
+                db_session.query(AvailableTenant)
+                .order_by(AvailableTenant.date_created)
+                .with_for_update(skip_locked=True)  # Skip locked rows to avoid blocking
+                .first()
+            )
+
+            if available_tenant:
+                tenant_id = available_tenant.tenant_id
+                # Remove the tenant from the available tenants table
+                db_session.delete(available_tenant)
+                db_session.commit()
+                logger.info(f"Using pre-provisioned tenant {tenant_id}")
+                return tenant_id
+            else:
+                db_session.rollback()
+                return None
+        except Exception:
+            logger.exception("Error getting available tenant")
+            db_session.rollback()
+            return None
+
+
+async def setup_tenant(tenant_id: str) -> None:
+    """
+    Set up a tenant with all necessary configurations.
+    This is a centralized function that handles all tenant setup logic.
+    """
+    token = None
+    try:
+        token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+
+        # Run Alembic migrations
+        await asyncio.to_thread(run_alembic_migrations, tenant_id)
+
+        # Configure the tenant with default settings
+        with get_session_with_tenant(tenant_id=tenant_id) as db_session:
+            # Configure default API keys
+            configure_default_api_keys(db_session)
+
+            # Set up Onyx with appropriate settings
+            current_search_settings = (
+                db_session.query(SearchSettings)
+                .filter_by(status=IndexModelStatus.FUTURE)
+                .first()
+            )
+            cohere_enabled = (
+                current_search_settings is not None
+                and current_search_settings.provider_type == EmbeddingProvider.COHERE
+            )
+            setup_onyx(db_session, tenant_id, cohere_enabled=cohere_enabled)
+
+    except Exception as e:
+        logger.exception(f"Failed to set up tenant {tenant_id}")
+        raise e
+    finally:
+        if token is not None:
+            CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
+
+
+async def assign_tenant_to_user(
+    tenant_id: str, email: str, referral_source: str | None = None
+) -> None:
+    """
+    Assign a tenant to a user and perform necessary operations.
+    Uses transaction handling to ensure atomicity and includes retry logic
+    for control plane notifications.
+    """
+    # First, add the user to the tenant in a transaction
+
+    try:
+        add_users_to_tenant([email], tenant_id)
+
+        # Create milestone record in the same transaction context as the tenant assignment
+        with get_session_with_tenant(tenant_id=tenant_id) as db_session:
+            create_milestone_and_report(
+                user=None,
+                distinct_id=tenant_id,
+                event_type=MilestoneRecordType.TENANT_CREATED,
+                properties={
+                    "email": email,
+                },
+                db_session=db_session,
+            )
+    except Exception:
+        logger.exception(f"Failed to assign tenant {tenant_id} to user {email}")
+        raise Exception("Failed to assign tenant to user")
+
+    # Notify control plane with retry logic
+    if not DEV_MODE:
+        await notify_control_plane(tenant_id, email, referral_source)

backend/ee/onyx/server/tenants/schema_management.py

@@ -74,3 +74,21 @@ def drop_schema(tenant_id: str) -> None:
             text("DROP SCHEMA IF EXISTS %(schema_name)s CASCADE"),
             {"schema_name": tenant_id},
         )
+
+
+def get_current_alembic_version(tenant_id: str) -> str:
+    """Get the current Alembic version for a tenant."""
+    from alembic.runtime.migration import MigrationContext
+    from sqlalchemy import text
+
+    engine = get_sqlalchemy_engine()
+
+    # Set the search path to the tenant's schema
+    with engine.connect() as connection:
+        connection.execute(text(f'SET search_path TO "{tenant_id}"'))
+
+        # Get the current version from the alembic_version table
+        context = MigrationContext.configure(connection)
+        current_rev = context.get_current_revision()
+
+    return current_rev or "head"

backend/ee/onyx/server/tenants/user_mapping.py

@@ -67,15 +67,39 @@ def user_owns_a_tenant(email: str) -> bool:
 
 
 def add_users_to_tenant(emails: list[str], tenant_id: str) -> None:
+    """
+    Add users to a tenant with proper transaction handling.
+    Checks if users already have a tenant mapping to avoid duplicates.
+    """
     with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session:
         try:
+            # Start a transaction
+            db_session.begin()
+
             for email in emails:
-                db_session.add(
-                    UserTenantMapping(email=email, tenant_id=tenant_id, active=False)
+                # Check if the user already has a mapping to this tenant
+                existing_mapping = (
+                    db_session.query(UserTenantMapping)
+                    .filter(
+                        UserTenantMapping.email == email,
+                        UserTenantMapping.tenant_id == tenant_id,
+                    )
+                    .with_for_update()
+                    .first()
                 )
+
+                if not existing_mapping:
+                    # Only add if mapping doesn't exist
+                    db_session.add(UserTenantMapping(email=email, tenant_id=tenant_id))
+
+            # Commit the transaction
+            db_session.commit()
+            logger.info(f"Successfully added users {emails} to tenant {tenant_id}")
+
         except Exception:
             logger.exception(f"Failed to add users to tenant {tenant_id}")
-        db_session.commit()
+            db_session.rollback()
+            raise
 
 
 def remove_users_from_tenant(emails: list[str], tenant_id: str) -> None:

backend/onyx/background/celery/apps/light.py

@@ -112,5 +112,6 @@ celery_app.autodiscover_tasks(
         "onyx.background.celery.tasks.connector_deletion",
         "onyx.background.celery.tasks.doc_permission_syncing",
         "onyx.background.celery.tasks.indexing",
+        "onyx.background.celery.tasks.tenant_provisioning",
     ]
 )

backend/onyx/background/celery/apps/monitoring.py

@@ -92,5 +92,6 @@ def on_setup_logging(
 celery_app.autodiscover_tasks(
     [
         "onyx.background.celery.tasks.monitoring",
+        "onyx.background.celery.tasks.tenant_provisioning",
     ]
 )

backend/onyx/background/celery/tasks/beat_schedule.py

@@ -167,6 +167,16 @@ beat_cloud_tasks: list[dict] = [
             "expires": BEAT_EXPIRES_DEFAULT,
         },
     },
+    {
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-available-tenants",
+        "task": OnyxCeleryTask.CHECK_AVAILABLE_TENANTS,
+        "schedule": timedelta(minutes=10),
+        "options": {
+            "queue": OnyxCeleryQueues.MONITORING,
+            "priority": OnyxCeleryPriority.HIGH,
+            "expires": BEAT_EXPIRES_DEFAULT,
+        },
+    },
 ]
 
 # tasks that only run self hosted

backend/onyx/background/celery/tasks/tenant_provisioning/tasks.py

@@ -0,0 +1,199 @@
+"""
+Periodic tasks for tenant pre-provisioning.
+"""
+import asyncio
+import datetime
+import uuid
+
+from celery import shared_task
+from celery import Task
+from redis.lock import Lock as RedisLock
+
+from ee.onyx.server.tenants.provisioning import setup_tenant
+from ee.onyx.server.tenants.schema_management import create_schema_if_not_exists
+from ee.onyx.server.tenants.schema_management import get_current_alembic_version
+from onyx.background.celery.apps.app_base import task_logger
+from onyx.configs.app_configs import JOB_TIMEOUT
+from onyx.configs.app_configs import TARGET_AVAILABLE_TENANTS
+from onyx.configs.constants import OnyxCeleryPriority
+from onyx.configs.constants import OnyxCeleryQueues
+from onyx.configs.constants import OnyxCeleryTask
+from onyx.configs.constants import OnyxRedisLocks
+from onyx.db.engine import get_session_with_shared_schema
+from onyx.db.models import AvailableTenant
+from onyx.redis.redis_pool import get_redis_client
+from shared_configs.configs import MULTI_TENANT
+from shared_configs.configs import TENANT_ID_PREFIX
+
+# Default number of pre-provisioned tenants to maintain
+DEFAULT_TARGET_AVAILABLE_TENANTS = 5
+
+# Soft time limit for tenant pre-provisioning tasks (in seconds)
+_TENANT_PROVISIONING_SOFT_TIME_LIMIT = 60 * 5  # 5 minutes
+# Hard time limit for tenant pre-provisioning tasks (in seconds)
+_TENANT_PROVISIONING_TIME_LIMIT = 60 * 10  # 10 minutes
+
+
+@shared_task(
+    name=OnyxCeleryTask.CHECK_AVAILABLE_TENANTS,
+    queue=OnyxCeleryQueues.MONITORING,
+    ignore_result=True,
+    soft_time_limit=JOB_TIMEOUT,
+    trail=False,
+    bind=True,
+)
+def check_available_tenants(self: Task) -> None:
+    """
+    Check if we have enough pre-provisioned tenants available.
+    If not, trigger the pre-provisioning of new tenants.
+    """
+    task_logger.info("STARTING CHECK_AVAILABLE_TENANTS")
+    if not MULTI_TENANT:
+        task_logger.info(
+            "Multi-tenancy is not enabled, skipping tenant pre-provisioning"
+        )
+        return
+
+    r = get_redis_client()
+    lock_check: RedisLock = r.lock(
+        OnyxRedisLocks.CHECK_AVAILABLE_TENANTS_LOCK,
+        timeout=_TENANT_PROVISIONING_SOFT_TIME_LIMIT,
+    )
+
+    # These tasks should never overlap
+    if not lock_check.acquire(blocking=False):
+        task_logger.info(
+            "Skipping check_available_tenants task because it is already running"
+        )
+        return
+
+    try:
+        # Get the current count of available tenants
+        with get_session_with_shared_schema() as db_session:
+            available_tenants_count = db_session.query(AvailableTenant).count()
+
+        # Get the target number of available tenants
+        target_available_tenants = getattr(
+            TARGET_AVAILABLE_TENANTS, "value", DEFAULT_TARGET_AVAILABLE_TENANTS
+        )
+
+        # Calculate how many new tenants we need to provision
+        tenants_to_provision = max(
+            0, target_available_tenants - available_tenants_count
+        )
+
+        task_logger.info(
+            f"Available tenants: {available_tenants_count}, "
+            f"Target: {target_available_tenants}, "
+            f"To provision: {tenants_to_provision}"
+        )
+
+        # Trigger pre-provisioning tasks for each tenant needed
+        for _ in range(tenants_to_provision):
+            from celery import current_app
+
+            current_app.send_task(
+                OnyxCeleryTask.PRE_PROVISION_TENANT,
+                priority=OnyxCeleryPriority.LOW,
+            )
+
+    except Exception:
+        task_logger.exception("Error in check_available_tenants task")
+
+    finally:
+        lock_check.release()
+
+
+@shared_task(
+    name=OnyxCeleryTask.PRE_PROVISION_TENANT,
+    ignore_result=True,
+    soft_time_limit=_TENANT_PROVISIONING_SOFT_TIME_LIMIT,
+    time_limit=_TENANT_PROVISIONING_TIME_LIMIT,
+    queue=OnyxCeleryQueues.MONITORING,
+    bind=True,
+)
+def pre_provision_tenant(self: Task) -> None:
+    """
+    Pre-provision a new tenant and store it in the NewAvailableTenant table.
+    This function fully sets up the tenant with all necessary configurations,
+    so it's ready to be assigned to a user immediately.
+    """
+    # The MULTI_TENANT check is now done at the caller level (check_available_tenants)
+    # rather than inside this function
+
+    r = get_redis_client()
+    lock_provision: RedisLock = r.lock(
+        OnyxRedisLocks.PRE_PROVISION_TENANT_LOCK,
+        timeout=_TENANT_PROVISIONING_SOFT_TIME_LIMIT,
+    )
+
+    # Allow multiple pre-provisioning tasks to run, but ensure they don't overlap
+    if not lock_provision.acquire(blocking=False):
+        task_logger.debug(
+            "Skipping pre_provision_tenant task because it is already running"
+        )
+        return
+
+    tenant_id: str | None = None
+    try:
+        # Generate a new tenant ID
+        tenant_id = TENANT_ID_PREFIX + str(uuid.uuid4())
+        task_logger.info(f"Pre-provisioning tenant: {tenant_id}")
+
+        # Create the schema for the new tenant
+        schema_created = create_schema_if_not_exists(tenant_id)
+        if schema_created:
+            task_logger.debug(f"Created schema for tenant: {tenant_id}")
+        else:
+            task_logger.debug(f"Schema already exists for tenant: {tenant_id}")
+
+        # Set up the tenant with all necessary configurations
+        task_logger.debug(f"Setting up tenant configuration: {tenant_id}")
+        asyncio.run(setup_tenant(tenant_id))
+        task_logger.debug(f"Tenant configuration completed: {tenant_id}")
+
+        # Get the current Alembic version
+        alembic_version = get_current_alembic_version(tenant_id)
+        task_logger.debug(
+            f"Tenant {tenant_id} using Alembic version: {alembic_version}"
+        )
+
+        # Store the pre-provisioned tenant in the database
+        task_logger.debug(f"Storing pre-provisioned tenant in database: {tenant_id}")
+        with get_session_with_shared_schema() as db_session:
+            # Use a transaction to ensure atomicity
+            db_session.begin()
+            try:
+                new_tenant = AvailableTenant(
+                    tenant_id=tenant_id,
+                    alembic_version=alembic_version,
+                    date_created=datetime.datetime.now(),
+                )
+                db_session.add(new_tenant)
+                db_session.commit()
+                task_logger.info(f"Successfully pre-provisioned tenant: {tenant_id}")
+            except Exception:
+                db_session.rollback()
+                task_logger.error(
+                    f"Failed to store pre-provisioned tenant: {tenant_id}",
+                    exc_info=True,
+                )
+                raise
+
+    except Exception:
+        task_logger.error("Error in pre_provision_tenant task", exc_info=True)
+        # If we have a tenant_id, attempt to rollback any partially completed provisioning
+        if tenant_id:
+            task_logger.info(
+                f"Rolling back failed tenant provisioning for: {tenant_id}"
+            )
+            try:
+                from ee.onyx.server.tenants.provisioning import (
+                    rollback_tenant_provisioning,
+                )
+
+                asyncio.run(rollback_tenant_provisioning(tenant_id))
+            except Exception:
+                task_logger.exception(f"Error during rollback for tenant: {tenant_id}")
+    finally:
+        lock_provision.release()

backend/onyx/configs/app_configs.py

@@ -643,3 +643,6 @@ MOCK_LLM_RESPONSE = (
 
 
 DEFAULT_IMAGE_ANALYSIS_MAX_SIZE_MB = 20
+
+# Number of pre-provisioned tenants to maintain
+TARGET_AVAILABLE_TENANTS = int(os.environ.get("TARGET_AVAILABLE_TENANTS", "5"))

backend/onyx/configs/constants.py

@@ -322,6 +322,8 @@ class OnyxRedisLocks:
         "da_lock:check_connector_external_group_sync_beat"
     )
     MONITOR_BACKGROUND_PROCESSES_LOCK = "da_lock:monitor_background_processes"
+    CHECK_AVAILABLE_TENANTS_LOCK = "da_lock:check_available_tenants"
+    PRE_PROVISION_TENANT_LOCK = "da_lock:pre_provision_tenant"
 
     CONNECTOR_DOC_PERMISSIONS_SYNC_LOCK_PREFIX = (
         "da_lock:connector_doc_permissions_sync"
@@ -384,6 +386,7 @@ class OnyxCeleryTask:
     CLOUD_MONITOR_CELERY_QUEUES = (
         f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor_celery_queues"
     )
+    CHECK_AVAILABLE_TENANTS = f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check_available_tenants"
 
     CHECK_FOR_CONNECTOR_DELETION = "check_for_connector_deletion_task"
     CHECK_FOR_VESPA_SYNC_TASK = "check_for_vespa_sync_task"
@@ -400,6 +403,9 @@ class OnyxCeleryTask:
     MONITOR_BACKGROUND_PROCESSES = "monitor_background_processes"
     MONITOR_CELERY_QUEUES = "monitor_celery_queues"
 
+    # Tenant pre-provisioning
+    PRE_PROVISION_TENANT = "pre_provision_tenant"
+
     KOMBU_MESSAGE_CLEANUP_TASK = "kombu_message_cleanup_task"
     CONNECTOR_PERMISSION_SYNC_GENERATOR_TASK = (
         "connector_permission_sync_generator_task"

backend/onyx/connectors/confluence/connector.py

@@ -263,6 +263,7 @@ class ConfluenceConnector(
                     result = process_attachment(
                         self.confluence_client,
                         attachment,
+                        page_id,
                         page_title,
                         self.image_analysis_llm,
                     )
@@ -366,6 +367,7 @@ class ConfluenceConnector(
                     response = convert_attachment_to_content(
                         confluence_client=self.confluence_client,
                         attachment=attachment,
+                        page_id=page["id"],
                         page_context=confluence_xml,
                         llm=self.image_analysis_llm,
                     )

backend/onyx/connectors/confluence/onyx_confluence.py

@@ -1,4 +1,3 @@
-import io
 import json
 import time
 from collections.abc import Callable
@@ -19,17 +18,11 @@ from requests import HTTPError
 
 from ee.onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_ID
 from ee.onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET
-from onyx.configs.app_configs import (
-    CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD,
-)
-from onyx.configs.app_configs import CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD
 from onyx.connectors.confluence.utils import _handle_http_error
 from onyx.connectors.confluence.utils import confluence_refresh_tokens
 from onyx.connectors.confluence.utils import get_start_param_from_url
 from onyx.connectors.confluence.utils import update_param_in_path
-from onyx.connectors.confluence.utils import validate_attachment_filetype
 from onyx.connectors.interfaces import CredentialsProviderInterface
-from onyx.file_processing.extract_file_text import extract_file_text
 from onyx.file_processing.html_utils import format_document_soup
 from onyx.redis.redis_pool import get_redis_client
 from onyx.utils.logger import setup_logger
@@ -808,65 +801,6 @@ def _get_user(confluence_client: OnyxConfluence, user_id: str) -> str:
     return _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) or _USER_NOT_FOUND
 
 
-def attachment_to_content(
-    confluence_client: OnyxConfluence,
-    attachment: dict[str, Any],
-    parent_content_id: str | None = None,
-) -> str | None:
-    """If it returns None, assume that we should skip this attachment."""
-    if not validate_attachment_filetype(attachment):
-        return None
-
-    if "api.atlassian.com" in confluence_client.url:
-        # https://developer.atlassian.com/cloud/confluence/rest/v1/api-group-content---attachments/#api-wiki-rest-api-content-id-child-attachment-attachmentid-download-get
-        if not parent_content_id:
-            logger.warning(
-                "parent_content_id is required to download attachments from Confluence Cloud!"
-            )
-            return None
-
-        download_link = (
-            confluence_client.url
-            + f"/rest/api/content/{parent_content_id}/child/attachment/{attachment['id']}/download"
-        )
-    else:
-        download_link = confluence_client.url + attachment["_links"]["download"]
-
-    attachment_size = attachment["extensions"]["fileSize"]
-    if attachment_size > CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD:
-        logger.warning(
-            f"Skipping {download_link} due to size. "
-            f"size={attachment_size} "
-            f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD}"
-        )
-        return None
-
-    logger.info(f"_attachment_to_content - _session.get: link={download_link}")
-
-    # why are we using session.get here? we probably won't retry these ... is that ok?
-    response = confluence_client._session.get(download_link)
-    if response.status_code != 200:
-        logger.warning(
-            f"Failed to fetch {download_link} with invalid status code {response.status_code}"
-        )
-        return None
-
-    extracted_text = extract_file_text(
-        io.BytesIO(response.content),
-        file_name=attachment["title"],
-        break_on_unprocessable=False,
-    )
-    if len(extracted_text) > CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD:
-        logger.warning(
-            f"Skipping {download_link} due to char count. "
-            f"char count={len(extracted_text)} "
-            f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD}"
-        )
-        return None
-
-    return extracted_text
-
-
 def extract_text_from_confluence_html(
     confluence_client: OnyxConfluence,
     confluence_object: dict[str, Any],

backend/onyx/connectors/confluence/utils.py

@@ -22,6 +22,7 @@ from sqlalchemy.orm import Session
 from onyx.configs.app_configs import (
     CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD,
 )
+from onyx.configs.app_configs import CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD
 from onyx.configs.constants import FileOrigin
 
 if TYPE_CHECKING:
@@ -84,25 +85,35 @@ class AttachmentProcessingResult(BaseModel):
     error: str | None = None
 
 
-def _download_attachment(
-    confluence_client: "OnyxConfluence", attachment: dict[str, Any]
-) -> bytes | None:
-    """
-    Retrieves the raw bytes of an attachment from Confluence. Returns None on error.
-    """
-    download_link = confluence_client.url + attachment["_links"]["download"]
-    resp = confluence_client._session.get(download_link)
-    if resp.status_code != 200:
-        logger.warning(
-            f"Failed to fetch {download_link} with status code {resp.status_code}"
+def _make_attachment_link(
+    confluence_client: "OnyxConfluence",
+    attachment: dict[str, Any],
+    parent_content_id: str | None = None,
+) -> str | None:
+    download_link = ""
+
+    if "api.atlassian.com" in confluence_client.url:
+        # https://developer.atlassian.com/cloud/confluence/rest/v1/api-group-content---attachments/#api-wiki-rest-api-content-id-child-attachment-attachmentid-download-get
+        if not parent_content_id:
+            logger.warning(
+                "parent_content_id is required to download attachments from Confluence Cloud!"
+            )
+            return None
+
+        download_link = (
+            confluence_client.url
+            + f"/rest/api/content/{parent_content_id}/child/attachment/{attachment['id']}/download"
         )
-        return None
-    return resp.content
+    else:
+        download_link = confluence_client.url + attachment["_links"]["download"]
+
+    return download_link
 
 
 def process_attachment(
     confluence_client: "OnyxConfluence",
     attachment: dict[str, Any],
+    parent_content_id: str | None,
     page_context: str,
     llm: LLM | None,
 ) -> AttachmentProcessingResult:
@@ -122,11 +133,52 @@ def process_attachment(
                 error=f"Unsupported file type: {media_type}",
             )
 
-        # Download the attachment
-        raw_bytes = _download_attachment(confluence_client, attachment)
-        if raw_bytes is None:
+        attachment_link = _make_attachment_link(
+            confluence_client, attachment, parent_content_id
+        )
+        if not attachment_link:
             return AttachmentProcessingResult(
-                text=None, file_name=None, error="Failed to download attachment"
+                text=None, file_name=None, error="Failed to make attachment link"
+            )
+
+        attachment_size = attachment["extensions"]["fileSize"]
+
+        if not media_type.startswith("image/") or not llm:
+            if attachment_size > CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD:
+                logger.warning(
+                    f"Skipping {attachment_link} due to size. "
+                    f"size={attachment_size} "
+                    f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD}"
+                )
+                return AttachmentProcessingResult(
+                    text=None,
+                    file_name=None,
+                    error=f"Attachment text too long: {attachment_size} chars",
+                )
+
+        logger.info(
+            f"Downloading attachment: "
+            f"title={attachment['title']} "
+            f"length={attachment_size} "
+            f"link={attachment_link}"
+        )
+
+        # Download the attachment
+        resp: requests.Response = confluence_client._session.get(attachment_link)
+        if resp.status_code != 200:
+            logger.warning(
+                f"Failed to fetch {attachment_link} with status code {resp.status_code}"
+            )
+            return AttachmentProcessingResult(
+                text=None,
+                file_name=None,
+                error=f"Attachment download status code is {resp.status_code}",
+            )
+
+        raw_bytes = resp.content
+        if not raw_bytes:
+            return AttachmentProcessingResult(
+                text=None, file_name=None, error="attachment.content is None"
             )
 
         # Process image attachments with LLM if available
@@ -249,6 +301,7 @@ def _process_text_attachment(
 def convert_attachment_to_content(
     confluence_client: "OnyxConfluence",
     attachment: dict[str, Any],
+    page_id: str,
     page_context: str,
     llm: LLM | None,
 ) -> tuple[str | None, str | None] | None:
@@ -266,7 +319,9 @@ def convert_attachment_to_content(
         )
         return None
 
-    result = process_attachment(confluence_client, attachment, page_context, llm)
+    result = process_attachment(
+        confluence_client, attachment, page_id, page_context, llm
+    )
     if result.error is not None:
         logger.warning(
             f"Attachment {attachment['title']} encountered error: {result.error}"

backend/onyx/connectors/salesforce/connector.py

@@ -48,10 +48,12 @@ class SalesforceConnector(LoadConnector, PollConnector, SlimConnector):
         self,
         credentials: dict[str, Any],
     ) -> dict[str, Any] | None:
+        domain = "test" if credentials.get("is_sandbox") else None
         self._sf_client = Salesforce(
             username=credentials["sf_username"],
             password=credentials["sf_password"],
             security_token=credentials["sf_security_token"],
+            domain=domain,
         )
         return None
 

backend/onyx/db/models.py

@@ -2309,6 +2309,17 @@ class UserTenantMapping(Base):
         return value.lower() if value else value
 
 
+class AvailableTenant(Base):
+    __tablename__ = "available_tenant"
+    """
+    These entries will only exist ephemerally and are meant to be picked up by new users on registration.
+    """
+
+    tenant_id: Mapped[str] = mapped_column(String, primary_key=True, nullable=False)
+    alembic_version: Mapped[str] = mapped_column(String, nullable=False)
+    date_created: Mapped[datetime.datetime] = mapped_column(DateTime, nullable=False)
+
+
 # This is a mapping from tenant IDs to anonymous user paths
 class TenantAnonymousUserPath(Base):
     __tablename__ = "tenant_anonymous_user_path"

web/src/app/admin/connectors/[connector]/pages/gdrive/Credential.tsx

@@ -1,4 +1,3 @@
-import { Button } from "@/components/Button";
 import { PopupSpec } from "@/components/admin/connectors/Popup";
 import React, { useState, useEffect } from "react";
 import { useSWRConfig } from "swr";
@@ -8,10 +7,14 @@ import { adminDeleteCredential } from "@/lib/credential";
 import { setupGoogleDriveOAuth } from "@/lib/googleDrive";
 import { GOOGLE_DRIVE_AUTH_IS_ADMIN_COOKIE_NAME } from "@/lib/constants";
 import Cookies from "js-cookie";
-import { TextFormField } from "@/components/admin/connectors/Field";
+import {
+  TextFormField,
+  SectionHeader,
+  SubLabel,
+} from "@/components/admin/connectors/Field";
 import { Form, Formik } from "formik";
 import { User } from "@/lib/types";
-import { Button as TremorButton } from "@/components/ui/button";
+import { Button } from "@/components/ui/button";
 import {
   Credential,
   GoogleDriveCredentialJson,
@@ -20,6 +23,15 @@ import {
 import { refreshAllGoogleData } from "@/lib/googleConnector";
 import { ValidSources } from "@/lib/types";
 import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib";
+import {
+  FiFile,
+  FiUpload,
+  FiTrash2,
+  FiCheck,
+  FiLink,
+  FiAlertTriangle,
+} from "react-icons/fi";
+import { cn, truncateString } from "@/lib/utils";
 
 type GoogleDriveCredentialJsonTypes = "authorized_user" | "service_account";
 
@@ -31,126 +43,202 @@ export const DriveJsonUpload = ({
   onSuccess?: () => void;
 }) => {
   const { mutate } = useSWRConfig();
-  const [credentialJsonStr, setCredentialJsonStr] = useState<
-    string | undefined
-  >();
+  const [isUploading, setIsUploading] = useState(false);
+  const [fileName, setFileName] = useState<string | undefined>();
+  const [isDragging, setIsDragging] = useState(false);
+
+  const handleFileUpload = async (file: File) => {
+    setIsUploading(true);
+    setFileName(file.name);
+
+    const reader = new FileReader();
+    reader.onload = async (loadEvent) => {
+      if (!loadEvent?.target?.result) {
+        setIsUploading(false);
+        return;
+      }
+
+      const credentialJsonStr = loadEvent.target.result as string;
+
+      // Check credential type
+      let credentialFileType: GoogleDriveCredentialJsonTypes;
+      try {
+        const appCredentialJson = JSON.parse(credentialJsonStr);
+        if (appCredentialJson.web) {
+          credentialFileType = "authorized_user";
+        } else if (appCredentialJson.type === "service_account") {
+          credentialFileType = "service_account";
+        } else {
+          throw new Error(
+            "Unknown credential type, expected one of 'OAuth Web application' or 'Service Account'"
+          );
+        }
+      } catch (e) {
+        setPopup({
+          message: `Invalid file provided - ${e}`,
+          type: "error",
+        });
+        setIsUploading(false);
+        return;
+      }
+
+      if (credentialFileType === "authorized_user") {
+        const response = await fetch(
+          "/api/manage/admin/connector/google-drive/app-credential",
+          {
+            method: "PUT",
+            headers: {
+              "Content-Type": "application/json",
+            },
+            body: credentialJsonStr,
+          }
+        );
+        if (response.ok) {
+          setPopup({
+            message: "Successfully uploaded app credentials",
+            type: "success",
+          });
+          mutate("/api/manage/admin/connector/google-drive/app-credential");
+          if (onSuccess) {
+            onSuccess();
+          }
+        } else {
+          const errorMsg = await response.text();
+          setPopup({
+            message: `Failed to upload app credentials - ${errorMsg}`,
+            type: "error",
+          });
+        }
+      }
+
+      if (credentialFileType === "service_account") {
+        const response = await fetch(
+          "/api/manage/admin/connector/google-drive/service-account-key",
+          {
+            method: "PUT",
+            headers: {
+              "Content-Type": "application/json",
+            },
+            body: credentialJsonStr,
+          }
+        );
+        if (response.ok) {
+          setPopup({
+            message: "Successfully uploaded service account key",
+            type: "success",
+          });
+          mutate(
+            "/api/manage/admin/connector/google-drive/service-account-key"
+          );
+          if (onSuccess) {
+            onSuccess();
+          }
+        } else {
+          const errorMsg = await response.text();
+          setPopup({
+            message: `Failed to upload service account key - ${errorMsg}`,
+            type: "error",
+          });
+        }
+      }
+      setIsUploading(false);
+    };
+
+    reader.readAsText(file);
+  };
+
+  const handleDragEnter = (e: React.DragEvent<HTMLLabelElement>) => {
+    e.preventDefault();
+    e.stopPropagation();
+    if (!isUploading) {
+      setIsDragging(true);
+    }
+  };
+
+  const handleDragLeave = (e: React.DragEvent<HTMLLabelElement>) => {
+    e.preventDefault();
+    e.stopPropagation();
+    setIsDragging(false);
+  };
+
+  const handleDragOver = (e: React.DragEvent<HTMLLabelElement>) => {
+    e.preventDefault();
+    e.stopPropagation();
+  };
+
+  const handleDrop = (e: React.DragEvent<HTMLLabelElement>) => {
+    e.preventDefault();
+    e.stopPropagation();
+    setIsDragging(false);
+
+    if (isUploading) return;
+
+    const files = e.dataTransfer.files;
+    if (files.length > 0) {
+      const file = files[0];
+      if (file.type === "application/json" || file.name.endsWith(".json")) {
+        handleFileUpload(file);
+      } else {
+        setPopup({
+          message: "Please upload a JSON file",
+          type: "error",
+        });
+      }
+    }
+  };
 
   return (
-    <>
-      <input
-        className={
-          "mr-3 text-sm text-text-900 border border-background-300 " +
-          "cursor-pointer bg-backgrournd dark:text-text-400 focus:outline-none " +
-          "dark:bg-background-700 dark:border-background-600 dark:placeholder-text-400"
-        }
-        type="file"
-        accept=".json"
-        onChange={(event) => {
-          if (!event.target.files) {
-            return;
-          }
-          const file = event.target.files[0];
-          const reader = new FileReader();
-
-          reader.onload = function (loadEvent) {
-            if (!loadEvent?.target?.result) {
-              return;
-            }
-            const fileContents = loadEvent.target.result;
-            setCredentialJsonStr(fileContents as string);
-          };
-
-          reader.readAsText(file);
-        }}
-      />
-
-      <Button
-        disabled={!credentialJsonStr}
-        onClick={async () => {
-          let credentialFileType: GoogleDriveCredentialJsonTypes;
-          try {
-            const appCredentialJson = JSON.parse(credentialJsonStr!);
-            if (appCredentialJson.web) {
-              credentialFileType = "authorized_user";
-            } else if (appCredentialJson.type === "service_account") {
-              credentialFileType = "service_account";
-            } else {
-              throw new Error(
-                "Unknown credential type, expected one of 'OAuth Web application' or 'Service Account'"
-              );
-            }
-          } catch (e) {
-            setPopup({
-              message: `Invalid file provided - ${e}`,
-              type: "error",
-            });
-            return;
-          }
-
-          if (credentialFileType === "authorized_user") {
-            const response = await fetch(
-              "/api/manage/admin/connector/google-drive/app-credential",
-              {
-                method: "PUT",
-                headers: {
-                  "Content-Type": "application/json",
-                },
-                body: credentialJsonStr,
-              }
-            );
-            if (response.ok) {
-              setPopup({
-                message: "Successfully uploaded app credentials",
-                type: "success",
-              });
-              mutate("/api/manage/admin/connector/google-drive/app-credential");
-              if (onSuccess) {
-                onSuccess();
-              }
-            } else {
-              const errorMsg = await response.text();
-              setPopup({
-                message: `Failed to upload app credentials - ${errorMsg}`,
-                type: "error",
-              });
-            }
-          }
-
-          if (credentialFileType === "service_account") {
-            const response = await fetch(
-              "/api/manage/admin/connector/google-drive/service-account-key",
-              {
-                method: "PUT",
-                headers: {
-                  "Content-Type": "application/json",
-                },
-                body: credentialJsonStr,
-              }
-            );
-            if (response.ok) {
-              setPopup({
-                message: "Successfully uploaded service account key",
-                type: "success",
-              });
-              mutate(
-                "/api/manage/admin/connector/google-drive/service-account-key"
-              );
-              if (onSuccess) {
-                onSuccess();
-              }
-            } else {
-              const errorMsg = await response.text();
-              setPopup({
-                message: `Failed to upload service account key - ${errorMsg}`,
-                type: "error",
-              });
-            }
-          }
-        }}
-      >
-        Upload
-      </Button>
-    </>
+    <div className="flex flex-col mt-4">
+      <div className="flex items-center">
+        <div className="relative flex flex-1 items-center">
+          <label
+            className={cn(
+              "flex h-10 items-center justify-center w-full px-4 py-2 border border-dashed rounded-md transition-colors",
+              isUploading
+                ? "opacity-70 cursor-not-allowed border-background-400 bg-background-50/30"
+                : isDragging
+                  ? "bg-background-50/50 border-primary dark:border-primary"
+                  : "cursor-pointer hover:bg-background-50/30 hover:border-primary dark:hover:border-primary border-background-300 dark:border-background-600"
+            )}
+            onDragEnter={handleDragEnter}
+            onDragLeave={handleDragLeave}
+            onDragOver={handleDragOver}
+            onDrop={handleDrop}
+          >
+            <div className="flex items-center space-x-2">
+              {isUploading ? (
+                <div className="h-4 w-4 border-t-2 border-b-2 border-primary rounded-full animate-spin"></div>
+              ) : (
+                <FiFile className="h-4 w-4 text-text-500" />
+              )}
+              <span className="text-sm text-text-500">
+                {isUploading
+                  ? `Uploading ${truncateString(fileName || "file", 50)}...`
+                  : isDragging
+                    ? "Drop JSON file here"
+                    : truncateString(
+                        fileName || "Select or drag JSON credentials file...",
+                        50
+                      )}
+              </span>
+            </div>
+            <input
+              className="sr-only"
+              type="file"
+              accept=".json"
+              disabled={isUploading}
+              onChange={(event) => {
+                if (!event.target.files?.length) {
+                  return;
+                }
+                const file = event.target.files[0];
+                handleFileUpload(file);
+              }}
+            />
+          </label>
+        </div>
+      </div>
+    </div>
   );
 };
 
@@ -160,6 +248,7 @@ interface DriveJsonUploadSectionProps {
   serviceAccountCredentialData?: { service_account_email: string };
   isAdmin: boolean;
   onSuccess?: () => void;
+  existingAuthCredential?: boolean;
 }
 
 export const DriveJsonUploadSection = ({
@@ -168,6 +257,7 @@ export const DriveJsonUploadSection = ({
   serviceAccountCredentialData,
   isAdmin,
   onSuccess,
+  existingAuthCredential,
 }: DriveJsonUploadSectionProps) => {
   const { mutate } = useSWRConfig();
   const router = useRouter();
@@ -177,6 +267,7 @@ export const DriveJsonUploadSection = ({
   const [localAppCredentialData, setLocalAppCredentialData] =
     useState(appCredentialData);
 
+  // Update local state when props change
   useEffect(() => {
     setLocalServiceAccountData(serviceAccountCredentialData);
     setLocalAppCredentialData(appCredentialData);
@@ -190,153 +281,135 @@ export const DriveJsonUploadSection = ({
     }
   };
 
-  if (localServiceAccountData?.service_account_email) {
-    return (
-      <div className="mt-2 text-sm">
-        <div>
-          Found existing service account key with the following <b>Email:</b>
-          <p className="italic mt-1">
-            {localServiceAccountData.service_account_email}
-          </p>
-        </div>
-        {isAdmin ? (
-          <>
-            <div className="mt-4 mb-1">
-              If you want to update these credentials, delete the existing
-              credentials through the button below, and then upload a new
-              credentials JSON.
-            </div>
-            <Button
-              onClick={async () => {
-                const response = await fetch(
-                  "/api/manage/admin/connector/google-drive/service-account-key",
-                  {
-                    method: "DELETE",
-                  }
-                );
-                if (response.ok) {
-                  mutate(
-                    "/api/manage/admin/connector/google-drive/service-account-key"
-                  );
-                  mutate(
-                    buildSimilarCredentialInfoURL(ValidSources.GoogleDrive)
-                  );
-                  setPopup({
-                    message: "Successfully deleted service account key",
-                    type: "success",
-                  });
-                  setLocalServiceAccountData(undefined);
-                  handleSuccess();
-                } else {
-                  const errorMsg = await response.text();
-                  setPopup({
-                    message: `Failed to delete service account key - ${errorMsg}`,
-                    type: "error",
-                  });
-                }
-              }}
-            >
-              Delete
-            </Button>
-          </>
-        ) : (
-          <>
-            <div className="mt-4 mb-1">
-              To change these credentials, please contact an administrator.
-            </div>
-          </>
-        )}
-      </div>
-    );
-  }
-
-  if (localAppCredentialData?.client_id) {
-    return (
-      <div className="mt-2 text-sm">
-        <div>
-          Found existing app credentials with the following <b>Client ID:</b>
-          <p className="italic mt-1">{localAppCredentialData.client_id}</p>
-        </div>
-        {isAdmin ? (
-          <>
-            <div className="mt-4 mb-1">
-              If you want to update these credentials, delete the existing
-              credentials through the button below, and then upload a new
-              credentials JSON.
-            </div>
-            <Button
-              onClick={async () => {
-                const response = await fetch(
-                  "/api/manage/admin/connector/google-drive/app-credential",
-                  {
-                    method: "DELETE",
-                  }
-                );
-                if (response.ok) {
-                  mutate(
-                    "/api/manage/admin/connector/google-drive/app-credential"
-                  );
-                  mutate(
-                    buildSimilarCredentialInfoURL(ValidSources.GoogleDrive)
-                  );
-                  setPopup({
-                    message: "Successfully deleted app credentials",
-                    type: "success",
-                  });
-                  setLocalAppCredentialData(undefined);
-                  handleSuccess();
-                } else {
-                  const errorMsg = await response.text();
-                  setPopup({
-                    message: `Failed to delete app credential - ${errorMsg}`,
-                    type: "error",
-                  });
-                }
-              }}
-            >
-              Delete
-            </Button>
-          </>
-        ) : (
-          <div className="mt-4 mb-1">
-            To change these credentials, please contact an administrator.
-          </div>
-        )}
-      </div>
-    );
-  }
-
   if (!isAdmin) {
     return (
-      <div className="mt-2">
-        <p className="text-sm mb-2">
-          Curators are unable to set up the google drive credentials. To add a
-          Google Drive connector, please contact an administrator.
-        </p>
+      <div>
+        <div className="flex items-start py-3 px-4 bg-yellow-50/30 dark:bg-yellow-900/5 rounded">
+          <FiAlertTriangle className="text-yellow-500 h-5 w-5 mr-2 mt-0.5 flex-shrink-0" />
+          <p className="text-sm">
+            Curators are unable to set up the Google Drive credentials. To add a
+            Google Drive connector, please contact an administrator.
+          </p>
+        </div>
       </div>
     );
   }
 
   return (
-    <div className="mt-2">
-      <p className="text-sm mb-2">
-        Follow the guide{" "}
+    <div>
+      <p className="text-sm mb-3">
+        To connect your Google Drive, create credentials (either OAuth App or
+        Service Account), download the JSON file, and upload it below.
+      </p>
+      <div className="mb-4">
         <a
-          className="text-link"
+          className="text-primary hover:text-primary/80 flex items-center gap-1 text-sm"
           target="_blank"
           href="https://docs.onyx.app/connectors/google_drive#authorization"
           rel="noreferrer"
         >
-          here
-        </a>{" "}
-        to either (1) setup a google OAuth App in your company workspace or (2)
-        create a Service Account.
-        <br />
-        <br />
-        Download the credentials JSON if choosing option (1) or the Service
-        Account key JSON if chooosing option (2), and upload it here.
-      </p>
-      <DriveJsonUpload setPopup={setPopup} onSuccess={handleSuccess} />
+          <FiLink className="h-3 w-3" />
+          View detailed setup instructions
+        </a>
+      </div>
+
+      {(localServiceAccountData?.service_account_email ||
+        localAppCredentialData?.client_id) && (
+        <div className="mb-4">
+          <div className="relative flex flex-1 items-center">
+            <label
+              className={cn(
+                "flex h-10 items-center justify-center w-full px-4 py-2 border border-dashed rounded-md transition-colors",
+                false
+                  ? "opacity-70 cursor-not-allowed border-background-400 bg-background-50/30"
+                  : "cursor-pointer hover:bg-background-50/30 hover:border-primary dark:hover:border-primary border-background-300 dark:border-background-600"
+              )}
+            >
+              <div className="flex items-center space-x-2">
+                {false ? (
+                  <div className="h-4 w-4 border-t-2 border-b-2 border-primary rounded-full animate-spin"></div>
+                ) : (
+                  <FiFile className="h-4 w-4 text-text-500" />
+                )}
+                <span className="text-sm text-text-500">
+                  {truncateString(
+                    localServiceAccountData?.service_account_email ||
+                      localAppCredentialData?.client_id ||
+                      "",
+                    50
+                  )}
+                </span>
+              </div>
+            </label>
+          </div>
+          {isAdmin && !existingAuthCredential && (
+            <div className="mt-2">
+              <Button
+                variant="destructive"
+                type="button"
+                onClick={async () => {
+                  const endpoint =
+                    localServiceAccountData?.service_account_email
+                      ? "/api/manage/admin/connector/google-drive/service-account-key"
+                      : "/api/manage/admin/connector/google-drive/app-credential";
+
+                  const response = await fetch(endpoint, {
+                    method: "DELETE",
+                  });
+
+                  if (response.ok) {
+                    mutate(endpoint);
+                    // Also mutate the credential endpoints to ensure Step 2 is reset
+                    mutate(
+                      buildSimilarCredentialInfoURL(ValidSources.GoogleDrive)
+                    );
+
+                    // Add additional mutations to refresh all credential-related endpoints
+                    mutate(
+                      "/api/manage/admin/connector/google-drive/credentials"
+                    );
+                    mutate(
+                      "/api/manage/admin/connector/google-drive/public-credential"
+                    );
+                    mutate(
+                      "/api/manage/admin/connector/google-drive/service-account-credential"
+                    );
+
+                    setPopup({
+                      message: `Successfully deleted ${
+                        localServiceAccountData
+                          ? "service account key"
+                          : "app credentials"
+                      }`,
+                      type: "success",
+                    });
+                    // Immediately update local state
+                    if (localServiceAccountData) {
+                      setLocalServiceAccountData(undefined);
+                    } else {
+                      setLocalAppCredentialData(undefined);
+                    }
+                    handleSuccess();
+                  } else {
+                    const errorMsg = await response.text();
+                    setPopup({
+                      message: `Failed to delete credentials - ${errorMsg}`,
+                      type: "error",
+                    });
+                  }
+                }}
+              >
+                Delete Credentials
+              </Button>
+            </div>
+          )}
+        </div>
+      )}
+
+      {!(
+        localServiceAccountData?.service_account_email ||
+        localAppCredentialData?.client_id
+      ) && <DriveJsonUpload setPopup={setPopup} onSuccess={handleSuccess} />}
     </div>
   );
 };
@@ -391,6 +464,7 @@ export const DriveAuthSection = ({
   user,
 }: DriveCredentialSectionProps) => {
   const router = useRouter();
+  const [isAuthenticating, setIsAuthenticating] = useState(false);
   const [localServiceAccountData, setLocalServiceAccountData] = useState(
     serviceAccountKeyData
   );
@@ -405,6 +479,7 @@ export const DriveAuthSection = ({
     setLocalGoogleDriveServiceAccountCredential,
   ] = useState(googleDriveServiceAccountCredential);
 
+  // Update local state when props change
   useEffect(() => {
     setLocalServiceAccountData(serviceAccountKeyData);
     setLocalAppCredentialData(appCredentialData);
@@ -424,126 +499,181 @@ export const DriveAuthSection = ({
     localGoogleDriveServiceAccountCredential;
   if (existingCredential) {
     return (
-      <>
-        <p className="mb-2 text-sm">
-          <i>Uploaded and authenticated credential already exists!</i>
-        </p>
-        <Button
-          onClick={async () => {
-            handleRevokeAccess(
-              connectorAssociated,
-              setPopup,
-              existingCredential,
-              refreshCredentials
-            );
-          }}
-        >
-          Revoke Access
-        </Button>
-      </>
+      <div>
+        <div className="mt-4">
+          <div className="py-3 px-4 bg-blue-50/30 dark:bg-blue-900/5 rounded mb-4 flex items-start">
+            <FiCheck className="text-blue-500 h-5 w-5 mr-2 mt-0.5 flex-shrink-0" />
+            <div className="flex-1">
+              <span className="font-medium block">Authentication Complete</span>
+              <p className="text-sm mt-1 text-text-500 dark:text-text-400 break-words">
+                Your Google Drive credentials have been successfully uploaded
+                and authenticated.
+              </p>
+            </div>
+          </div>
+          <Button
+            variant="destructive"
+            type="button"
+            onClick={async () => {
+              handleRevokeAccess(
+                connectorAssociated,
+                setPopup,
+                existingCredential,
+                refreshCredentials
+              );
+            }}
+          >
+            Revoke Access
+          </Button>
+        </div>
+      </div>
+    );
+  }
+
+  // If no credentials are uploaded, show message to complete step 1 first
+  if (
+    !localServiceAccountData?.service_account_email &&
+    !localAppCredentialData?.client_id
+  ) {
+    return (
+      <div>
+        <SectionHeader>Google Drive Authentication</SectionHeader>
+        <div className="mt-4">
+          <div className="flex items-start py-3 px-4 bg-yellow-50/30 dark:bg-yellow-900/5 rounded">
+            <FiAlertTriangle className="text-yellow-500 h-5 w-5 mr-2 mt-0.5 flex-shrink-0" />
+            <p className="text-sm">
+              Please complete Step 1 by uploading either OAuth credentials or a
+              Service Account key before proceeding with authentication.
+            </p>
+          </div>
+        </div>
+      </div>
     );
   }
 
   if (localServiceAccountData?.service_account_email) {
     return (
       <div>
-        <Formik
-          initialValues={{
-            google_primary_admin: user?.email || "",
-          }}
-          validationSchema={Yup.object().shape({
-            google_primary_admin: Yup.string().required(
-              "User email is required"
-            ),
-          })}
-          onSubmit={async (values, formikHelpers) => {
-            formikHelpers.setSubmitting(true);
-            const response = await fetch(
-              "/api/manage/admin/connector/google-drive/service-account-credential",
-              {
-                method: "PUT",
-                headers: {
-                  "Content-Type": "application/json",
-                },
-                body: JSON.stringify({
-                  google_primary_admin: values.google_primary_admin,
-                }),
-              }
-            );
+        <div className="mt-4">
+          <Formik
+            initialValues={{
+              google_primary_admin: user?.email || "",
+            }}
+            validationSchema={Yup.object().shape({
+              google_primary_admin: Yup.string()
+                .email("Must be a valid email")
+                .required("Required"),
+            })}
+            onSubmit={async (values, formikHelpers) => {
+              formikHelpers.setSubmitting(true);
+              try {
+                const response = await fetch(
+                  "/api/manage/admin/connector/google-drive/service-account-credential",
+                  {
+                    method: "PUT",
+                    headers: {
+                      "Content-Type": "application/json",
+                    },
+                    body: JSON.stringify({
+                      google_primary_admin: values.google_primary_admin,
+                    }),
+                  }
+                );
 
-            if (response.ok) {
-              setPopup({
-                message: "Successfully created service account credential",
-                type: "success",
-              });
-            } else {
-              const errorMsg = await response.text();
-              setPopup({
-                message: `Failed to create service account credential - ${errorMsg}`,
-                type: "error",
-              });
-            }
-            refreshCredentials();
-          }}
-        >
-          {({ isSubmitting }) => (
-            <Form>
-              <TextFormField
-                name="google_primary_admin"
-                label="Primary Admin Email:"
-                subtext="Enter the email of an admin/owner of the Google Organization that owns the Google Drive(s) you want to index."
-              />
-              <div className="flex">
-                <TremorButton type="submit" disabled={isSubmitting}>
-                  Create Credential
-                </TremorButton>
-              </div>
-            </Form>
-          )}
-        </Formik>
+                if (response.ok) {
+                  setPopup({
+                    message: "Successfully created service account credential",
+                    type: "success",
+                  });
+                  refreshCredentials();
+                } else {
+                  const errorMsg = await response.text();
+                  setPopup({
+                    message: `Failed to create service account credential - ${errorMsg}`,
+                    type: "error",
+                  });
+                }
+              } catch (error) {
+                setPopup({
+                  message: `Failed to create service account credential - ${error}`,
+                  type: "error",
+                });
+              } finally {
+                formikHelpers.setSubmitting(false);
+              }
+            }}
+          >
+            {({ isSubmitting }) => (
+              <Form>
+                <TextFormField
+                  name="google_primary_admin"
+                  label="Primary Admin Email:"
+                  subtext="Enter the email of an admin/owner of the Google Organization that owns the Google Drive(s) you want to index."
+                />
+                <div className="flex">
+                  <Button type="submit" disabled={isSubmitting}>
+                    {isSubmitting ? "Creating..." : "Create Credential"}
+                  </Button>
+                </div>
+              </Form>
+            )}
+          </Formik>
+        </div>
       </div>
     );
   }
 
   if (localAppCredentialData?.client_id) {
     return (
-      <div className="text-sm mb-4">
-        <p className="mb-2">
-          Next, you must provide credentials via OAuth. This gives us read
-          access to the docs you have access to in your google drive account.
-        </p>
+      <div>
+        <div className="bg-background-50/30 dark:bg-background-900/20 rounded mb-4">
+          <p className="text-sm">
+            Next, you need to authenticate with Google Drive via OAuth. This
+            gives us read access to the documents you have access to in your
+            Google Drive account.
+          </p>
+        </div>
         <Button
+          disabled={isAuthenticating}
           onClick={async () => {
-            const [authUrl, errorMsg] = await setupGoogleDriveOAuth({
-              isAdmin: true,
-              name: "OAuth (uploaded)",
-            });
-            if (authUrl) {
+            setIsAuthenticating(true);
+            try {
               // cookie used by callback to determine where to finally redirect to
               Cookies.set(GOOGLE_DRIVE_AUTH_IS_ADMIN_COOKIE_NAME, "true", {
                 path: "/",
               });
-              router.push(authUrl);
-              return;
-            }
 
-            setPopup({
-              message: errorMsg,
-              type: "error",
-            });
+              const [authUrl, errorMsg] = await setupGoogleDriveOAuth({
+                isAdmin: true,
+                name: "OAuth (uploaded)",
+              });
+
+              if (authUrl) {
+                router.push(authUrl);
+              } else {
+                setPopup({
+                  message: errorMsg,
+                  type: "error",
+                });
+                setIsAuthenticating(false);
+              }
+            } catch (error) {
+              setPopup({
+                message: `Failed to authenticate with Google Drive - ${error}`,
+                type: "error",
+              });
+              setIsAuthenticating(false);
+            }
           }}
         >
-          Authenticate with Google Drive
+          {isAuthenticating
+            ? "Authenticating..."
+            : "Authenticate with Google Drive"}
         </Button>
       </div>
     );
   }
 
-  // case where no keys have been uploaded in step 1
-  return (
-    <p className="text-sm">
-      Please upload either a OAuth Client Credential JSON or a Google Drive
-      Service Account Key JSON in Step 1 before moving onto Step 2.
-    </p>
-  );
+  // This code path should not be reached with the new conditions above
+  return null;
 };

web/src/app/admin/connectors/[connector]/pages/gdrive/GoogleDrivePage.tsx

@@ -165,6 +165,10 @@ const GDriveMain = ({
         serviceAccountCredentialData={serviceAccountKeyData}
         isAdmin={isAdmin}
         onSuccess={handleRefresh}
+        existingAuthCredential={Boolean(
+          googleDrivePublicUploadedCredential ||
+            googleDriveServiceAccountCredential
+        )}
       />
 
       {isAdmin &&

web/src/app/admin/connectors/[connector]/pages/gmail/Credential.tsx

@@ -1,4 +1,4 @@
-import { Button } from "@/components/Button";
+import { Button } from "@/components/ui/button";
 import { PopupSpec } from "@/components/admin/connectors/Popup";
 import React, { useState, useEffect } from "react";
 import { useSWRConfig } from "swr";
@@ -8,7 +8,11 @@ import { adminDeleteCredential } from "@/lib/credential";
 import { setupGmailOAuth } from "@/lib/gmail";
 import { GMAIL_AUTH_IS_ADMIN_COOKIE_NAME } from "@/lib/constants";
 import Cookies from "js-cookie";
-import { TextFormField } from "@/components/admin/connectors/Field";
+import {
+  TextFormField,
+  SectionHeader,
+  SubLabel,
+} from "@/components/admin/connectors/Field";
 import { Form, Formik } from "formik";
 import { User } from "@/lib/types";
 import CardSection from "@/components/admin/CardSection";
@@ -20,10 +24,19 @@ import {
 import { refreshAllGoogleData } from "@/lib/googleConnector";
 import { ValidSources } from "@/lib/types";
 import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib";
+import {
+  FiFile,
+  FiUpload,
+  FiTrash2,
+  FiCheck,
+  FiLink,
+  FiAlertTriangle,
+} from "react-icons/fi";
+import { cn, truncateString } from "@/lib/utils";
 
 type GmailCredentialJsonTypes = "authorized_user" | "service_account";
 
-const DriveJsonUpload = ({
+const GmailCredentialUpload = ({
   setPopup,
   onSuccess,
 }: {
@@ -31,134 +44,210 @@ const DriveJsonUpload = ({
   onSuccess?: () => void;
 }) => {
   const { mutate } = useSWRConfig();
-  const [credentialJsonStr, setCredentialJsonStr] = useState<
-    string | undefined
-  >();
+  const [isUploading, setIsUploading] = useState(false);
+  const [fileName, setFileName] = useState<string | undefined>();
+  const [isDragging, setIsDragging] = useState(false);
+
+  const handleFileUpload = async (file: File) => {
+    setIsUploading(true);
+    setFileName(file.name);
+
+    const reader = new FileReader();
+    reader.onload = async (loadEvent) => {
+      if (!loadEvent?.target?.result) {
+        setIsUploading(false);
+        return;
+      }
+
+      const credentialJsonStr = loadEvent.target.result as string;
+
+      // Check credential type
+      let credentialFileType: GmailCredentialJsonTypes;
+      try {
+        const appCredentialJson = JSON.parse(credentialJsonStr);
+        if (appCredentialJson.web) {
+          credentialFileType = "authorized_user";
+        } else if (appCredentialJson.type === "service_account") {
+          credentialFileType = "service_account";
+        } else {
+          throw new Error(
+            "Unknown credential type, expected one of 'OAuth Web application' or 'Service Account'"
+          );
+        }
+      } catch (e) {
+        setPopup({
+          message: `Invalid file provided - ${e}`,
+          type: "error",
+        });
+        setIsUploading(false);
+        return;
+      }
+
+      if (credentialFileType === "authorized_user") {
+        const response = await fetch(
+          "/api/manage/admin/connector/gmail/app-credential",
+          {
+            method: "PUT",
+            headers: {
+              "Content-Type": "application/json",
+            },
+            body: credentialJsonStr,
+          }
+        );
+        if (response.ok) {
+          setPopup({
+            message: "Successfully uploaded app credentials",
+            type: "success",
+          });
+          mutate("/api/manage/admin/connector/gmail/app-credential");
+          if (onSuccess) {
+            onSuccess();
+          }
+        } else {
+          const errorMsg = await response.text();
+          setPopup({
+            message: `Failed to upload app credentials - ${errorMsg}`,
+            type: "error",
+          });
+        }
+      }
+
+      if (credentialFileType === "service_account") {
+        const response = await fetch(
+          "/api/manage/admin/connector/gmail/service-account-key",
+          {
+            method: "PUT",
+            headers: {
+              "Content-Type": "application/json",
+            },
+            body: credentialJsonStr,
+          }
+        );
+        if (response.ok) {
+          setPopup({
+            message: "Successfully uploaded service account key",
+            type: "success",
+          });
+          mutate("/api/manage/admin/connector/gmail/service-account-key");
+          if (onSuccess) {
+            onSuccess();
+          }
+        } else {
+          const errorMsg = await response.text();
+          setPopup({
+            message: `Failed to upload service account key - ${errorMsg}`,
+            type: "error",
+          });
+        }
+      }
+      setIsUploading(false);
+    };
+
+    reader.readAsText(file);
+  };
+
+  const handleDragEnter = (e: React.DragEvent<HTMLLabelElement>) => {
+    e.preventDefault();
+    e.stopPropagation();
+    if (!isUploading) {
+      setIsDragging(true);
+    }
+  };
+
+  const handleDragLeave = (e: React.DragEvent<HTMLLabelElement>) => {
+    e.preventDefault();
+    e.stopPropagation();
+    setIsDragging(false);
+  };
+
+  const handleDragOver = (e: React.DragEvent<HTMLLabelElement>) => {
+    e.preventDefault();
+    e.stopPropagation();
+  };
+
+  const handleDrop = (e: React.DragEvent<HTMLLabelElement>) => {
+    e.preventDefault();
+    e.stopPropagation();
+    setIsDragging(false);
+
+    if (isUploading) return;
+
+    const files = e.dataTransfer.files;
+    if (files.length > 0) {
+      const file = files[0];
+      if (file.type === "application/json" || file.name.endsWith(".json")) {
+        handleFileUpload(file);
+      } else {
+        setPopup({
+          message: "Please upload a JSON file",
+          type: "error",
+        });
+      }
+    }
+  };
 
   return (
-    <>
-      <input
-        className={
-          "mr-3 text-sm text-text-900 border border-background-300 overflow-visible " +
-          "cursor-pointer bg-background dark:text-text-400 focus:outline-none " +
-          "dark:bg-background-700 dark:border-background-600 dark:placeholder-text-400"
-        }
-        type="file"
-        accept=".json"
-        onChange={(event) => {
-          if (!event.target.files) {
-            return;
-          }
-          const file = event.target.files[0];
-          const reader = new FileReader();
-
-          reader.onload = function (loadEvent) {
-            if (!loadEvent?.target?.result) {
-              return;
-            }
-            const fileContents = loadEvent.target.result;
-            setCredentialJsonStr(fileContents as string);
-          };
-
-          reader.readAsText(file);
-        }}
-      />
-
-      <Button
-        disabled={!credentialJsonStr}
-        onClick={async () => {
-          // check if the JSON is a app credential or a service account credential
-          let credentialFileType: GmailCredentialJsonTypes;
-          try {
-            const appCredentialJson = JSON.parse(credentialJsonStr!);
-            if (appCredentialJson.web) {
-              credentialFileType = "authorized_user";
-            } else if (appCredentialJson.type === "service_account") {
-              credentialFileType = "service_account";
-            } else {
-              throw new Error(
-                "Unknown credential type, expected one of 'OAuth Web application' or 'Service Account'"
-              );
-            }
-          } catch (e) {
-            setPopup({
-              message: `Invalid file provided - ${e}`,
-              type: "error",
-            });
-            return;
-          }
-
-          if (credentialFileType === "authorized_user") {
-            const response = await fetch(
-              "/api/manage/admin/connector/gmail/app-credential",
-              {
-                method: "PUT",
-                headers: {
-                  "Content-Type": "application/json",
-                },
-                body: credentialJsonStr,
-              }
-            );
-            if (response.ok) {
-              setPopup({
-                message: "Successfully uploaded app credentials",
-                type: "success",
-              });
-              mutate("/api/manage/admin/connector/gmail/app-credential");
-              if (onSuccess) {
-                onSuccess();
-              }
-            } else {
-              const errorMsg = await response.text();
-              setPopup({
-                message: `Failed to upload app credentials - ${errorMsg}`,
-                type: "error",
-              });
-            }
-          }
-
-          if (credentialFileType === "service_account") {
-            const response = await fetch(
-              "/api/manage/admin/connector/gmail/service-account-key",
-              {
-                method: "PUT",
-                headers: {
-                  "Content-Type": "application/json",
-                },
-                body: credentialJsonStr,
-              }
-            );
-            if (response.ok) {
-              setPopup({
-                message: "Successfully uploaded service account key",
-                type: "success",
-              });
-              mutate("/api/manage/admin/connector/gmail/service-account-key");
-              if (onSuccess) {
-                onSuccess();
-              }
-            } else {
-              const errorMsg = await response.text();
-              setPopup({
-                message: `Failed to upload service account key - ${errorMsg}`,
-                type: "error",
-              });
-            }
-          }
-        }}
-      >
-        Upload
-      </Button>
-    </>
+    <div className="flex flex-col mt-4">
+      <div className="flex items-center">
+        <div className="relative flex flex-1 items-center">
+          <label
+            className={cn(
+              "flex h-10 items-center justify-center w-full px-4 py-2 border border-dashed rounded-md transition-colors",
+              isUploading
+                ? "opacity-70 cursor-not-allowed border-background-400 bg-background-50/30"
+                : isDragging
+                  ? "bg-background-50/50 border-primary dark:border-primary"
+                  : "cursor-pointer hover:bg-background-50/30 hover:border-primary dark:hover:border-primary border-background-300 dark:border-background-600"
+            )}
+            onDragEnter={handleDragEnter}
+            onDragLeave={handleDragLeave}
+            onDragOver={handleDragOver}
+            onDrop={handleDrop}
+          >
+            <div className="flex items-center space-x-2">
+              {isUploading ? (
+                <div className="h-4 w-4 border-t-2 border-b-2 border-primary rounded-full animate-spin"></div>
+              ) : (
+                <FiFile className="h-4 w-4 text-text-500" />
+              )}
+              <span className="text-sm text-text-500">
+                {isUploading
+                  ? `Uploading ${truncateString(fileName || "file", 50)}...`
+                  : isDragging
+                    ? "Drop JSON file here"
+                    : truncateString(
+                        fileName || "Select or drag JSON credentials file...",
+                        50
+                      )}
+              </span>
+            </div>
+            <input
+              className="sr-only"
+              type="file"
+              accept=".json"
+              disabled={isUploading}
+              onChange={(event) => {
+                if (!event.target.files?.length) {
+                  return;
+                }
+                const file = event.target.files[0];
+                handleFileUpload(file);
+              }}
+            />
+          </label>
+        </div>
+      </div>
+    </div>
   );
 };
 
-interface DriveJsonUploadSectionProps {
+interface GmailJsonUploadSectionProps {
   setPopup: (popupSpec: PopupSpec | null) => void;
   appCredentialData?: { client_id: string };
   serviceAccountCredentialData?: { service_account_email: string };
   isAdmin: boolean;
   onSuccess?: () => void;
+  existingAuthCredential?: boolean;
 }
 
 export const GmailJsonUploadSection = ({
@@ -167,7 +256,8 @@ export const GmailJsonUploadSection = ({
   serviceAccountCredentialData,
   isAdmin,
   onSuccess,
-}: DriveJsonUploadSectionProps) => {
+  existingAuthCredential,
+}: GmailJsonUploadSectionProps) => {
   const { mutate } = useSWRConfig();
   const router = useRouter();
   const [localServiceAccountData, setLocalServiceAccountData] = useState(
@@ -190,156 +280,138 @@ export const GmailJsonUploadSection = ({
     }
   };
 
-  if (localServiceAccountData?.service_account_email) {
-    return (
-      <div className="mt-2 text-sm">
-        <div>
-          Found existing service account key with the following <b>Email:</b>
-          <p className="italic mt-1">
-            {localServiceAccountData.service_account_email}
-          </p>
-        </div>
-        {isAdmin ? (
-          <>
-            <div className="mt-4 mb-1">
-              If you want to update these credentials, delete the existing
-              credentials through the button below, and then upload a new
-              credentials JSON.
-            </div>
-            <Button
-              onClick={async () => {
-                const response = await fetch(
-                  "/api/manage/admin/connector/gmail/service-account-key",
-                  {
-                    method: "DELETE",
-                  }
-                );
-                if (response.ok) {
-                  mutate(
-                    "/api/manage/admin/connector/gmail/service-account-key"
-                  );
-                  // Also mutate the credential endpoints to ensure Step 2 is reset
-                  mutate(buildSimilarCredentialInfoURL(ValidSources.Gmail));
-                  setPopup({
-                    message: "Successfully deleted service account key",
-                    type: "success",
-                  });
-                  // Immediately update local state
-                  setLocalServiceAccountData(undefined);
-                  handleSuccess();
-                } else {
-                  const errorMsg = await response.text();
-                  setPopup({
-                    message: `Failed to delete service account key - ${errorMsg}`,
-                    type: "error",
-                  });
-                }
-              }}
-            >
-              Delete
-            </Button>
-          </>
-        ) : (
-          <>
-            <div className="mt-4 mb-1">
-              To change these credentials, please contact an administrator.
-            </div>
-          </>
-        )}
-      </div>
-    );
-  }
-
-  if (localAppCredentialData?.client_id) {
-    return (
-      <div className="mt-2 text-sm">
-        <div>
-          Found existing app credentials with the following <b>Client ID:</b>
-          <p className="italic mt-1">{localAppCredentialData.client_id}</p>
-        </div>
-        {isAdmin ? (
-          <>
-            <div className="mt-4 mb-1">
-              If you want to update these credentials, delete the existing
-              credentials through the button below, and then upload a new
-              credentials JSON.
-            </div>
-            <Button
-              onClick={async () => {
-                const response = await fetch(
-                  "/api/manage/admin/connector/gmail/app-credential",
-                  {
-                    method: "DELETE",
-                  }
-                );
-                if (response.ok) {
-                  mutate("/api/manage/admin/connector/gmail/app-credential");
-                  // Also mutate the credential endpoints to ensure Step 2 is reset
-                  mutate(buildSimilarCredentialInfoURL(ValidSources.Gmail));
-                  setPopup({
-                    message: "Successfully deleted app credentials",
-                    type: "success",
-                  });
-                  // Immediately update local state
-                  setLocalAppCredentialData(undefined);
-                  handleSuccess();
-                } else {
-                  const errorMsg = await response.text();
-                  setPopup({
-                    message: `Failed to delete app credential - ${errorMsg}`,
-                    type: "error",
-                  });
-                }
-              }}
-            >
-              Delete
-            </Button>
-          </>
-        ) : (
-          <div className="mt-4 mb-1">
-            To change these credentials, please contact an administrator.
-          </div>
-        )}
-      </div>
-    );
-  }
-
   if (!isAdmin) {
     return (
-      <div className="mt-2">
-        <p className="text-sm mb-2">
-          Curators are unable to set up the Gmail credentials. To add a Gmail
-          connector, please contact an administrator.
-        </p>
+      <div>
+        <div className="flex items-start py-3 px-4 bg-yellow-50/30 dark:bg-yellow-900/5 rounded">
+          <FiAlertTriangle className="text-yellow-500 h-5 w-5 mr-2 mt-0.5 flex-shrink-0" />
+          <p className="text-sm">
+            Curators are unable to set up the Gmail credentials. To add a Gmail
+            connector, please contact an administrator.
+          </p>
+        </div>
       </div>
     );
   }
 
   return (
-    <div className="mt-2">
-      <p className="text-sm mb-2">
-        Follow the guide{" "}
+    <div>
+      <p className="text-sm mb-3">
+        To connect your Gmail, create credentials (either OAuth App or Service
+        Account), download the JSON file, and upload it below.
+      </p>
+      <div className="mb-4">
         <a
-          className="text-link"
+          className="text-primary hover:text-primary/80 flex items-center gap-1 text-sm"
           target="_blank"
           href="https://docs.onyx.app/connectors/gmail#authorization"
           rel="noreferrer"
         >
-          here
-        </a>{" "}
-        to either (1) setup a Google OAuth App in your company workspace or (2)
-        create a Service Account.
-        <br />
-        <br />
-        Download the credentials JSON if choosing option (1) or the Service
-        Account key JSON if choosing option (2), and upload it here.
-      </p>
-      <DriveJsonUpload setPopup={setPopup} onSuccess={handleSuccess} />
+          <FiLink className="h-3 w-3" />
+          View detailed setup instructions
+        </a>
+      </div>
+
+      {(localServiceAccountData?.service_account_email ||
+        localAppCredentialData?.client_id) && (
+        <div className="mb-4">
+          <div className="relative flex flex-1 items-center">
+            <label
+              className={cn(
+                "flex h-10 items-center justify-center w-full px-4 py-2 border border-dashed rounded-md transition-colors",
+                false
+                  ? "opacity-70 cursor-not-allowed border-background-400 bg-background-50/30"
+                  : "cursor-pointer hover:bg-background-50/30 hover:border-primary dark:hover:border-primary border-background-300 dark:border-background-600"
+              )}
+            >
+              <div className="flex items-center space-x-2">
+                {false ? (
+                  <div className="h-4 w-4 border-t-2 border-b-2 border-primary rounded-full animate-spin"></div>
+                ) : (
+                  <FiFile className="h-4 w-4 text-text-500" />
+                )}
+                <span className="text-sm text-text-500">
+                  {truncateString(
+                    localServiceAccountData?.service_account_email ||
+                      localAppCredentialData?.client_id ||
+                      "",
+                    50
+                  )}
+                </span>
+              </div>
+            </label>
+          </div>
+          {isAdmin && !existingAuthCredential && (
+            <div className="mt-2">
+              <Button
+                variant="destructive"
+                type="button"
+                onClick={async () => {
+                  const endpoint =
+                    localServiceAccountData?.service_account_email
+                      ? "/api/manage/admin/connector/gmail/service-account-key"
+                      : "/api/manage/admin/connector/gmail/app-credential";
+
+                  const response = await fetch(endpoint, {
+                    method: "DELETE",
+                  });
+
+                  if (response.ok) {
+                    mutate(endpoint);
+                    // Also mutate the credential endpoints to ensure Step 2 is reset
+                    mutate(buildSimilarCredentialInfoURL(ValidSources.Gmail));
+
+                    // Add additional mutations to refresh all credential-related endpoints
+                    mutate("/api/manage/admin/connector/gmail/credentials");
+                    mutate(
+                      "/api/manage/admin/connector/gmail/public-credential"
+                    );
+                    mutate(
+                      "/api/manage/admin/connector/gmail/service-account-credential"
+                    );
+
+                    setPopup({
+                      message: `Successfully deleted ${
+                        localServiceAccountData
+                          ? "service account key"
+                          : "app credentials"
+                      }`,
+                      type: "success",
+                    });
+                    // Immediately update local state
+                    if (localServiceAccountData) {
+                      setLocalServiceAccountData(undefined);
+                    } else {
+                      setLocalAppCredentialData(undefined);
+                    }
+                    handleSuccess();
+                  } else {
+                    const errorMsg = await response.text();
+                    setPopup({
+                      message: `Failed to delete credentials - ${errorMsg}`,
+                      type: "error",
+                    });
+                  }
+                }}
+              >
+                Delete Credentials
+              </Button>
+            </div>
+          )}
+        </div>
+      )}
+
+      {!(
+        localServiceAccountData?.service_account_email ||
+        localAppCredentialData?.client_id
+      ) && (
+        <GmailCredentialUpload setPopup={setPopup} onSuccess={handleSuccess} />
+      )}
     </div>
   );
 };
 
-interface DriveCredentialSectionProps {
+interface GmailCredentialSectionProps {
   gmailPublicCredential?: Credential<GmailCredentialJson>;
   gmailServiceAccountCredential?: Credential<GmailServiceAccountCredentialJson>;
   serviceAccountKeyData?: { service_account_email: string };
@@ -387,7 +459,7 @@ export const GmailAuthSection = ({
   refreshCredentials,
   connectorExists,
   user,
-}: DriveCredentialSectionProps) => {
+}: GmailCredentialSectionProps) => {
   const router = useRouter();
   const [isAuthenticating, setIsAuthenticating] = useState(false);
   const [localServiceAccountData, setLocalServiceAccountData] = useState(
@@ -420,104 +492,141 @@ export const GmailAuthSection = ({
     localGmailPublicCredential || localGmailServiceAccountCredential;
   if (existingCredential) {
     return (
-      <>
-        <p className="mb-2 text-sm">
-          <i>Uploaded and authenticated credential already exists!</i>
-        </p>
-        <Button
-          onClick={async () => {
-            handleRevokeAccess(
-              connectorExists,
-              setPopup,
-              existingCredential,
-              refreshCredentials
-            );
-          }}
-        >
-          Revoke Access
-        </Button>
-      </>
+      <div>
+        <div className="mt-4">
+          <div className="py-3 px-4 bg-blue-50/30 dark:bg-blue-900/5 rounded mb-4 flex items-start">
+            <FiCheck className="text-blue-500 h-5 w-5 mr-2 mt-0.5 flex-shrink-0" />
+            <div className="flex-1">
+              <span className="font-medium block">Authentication Complete</span>
+              <p className="text-sm mt-1 text-text-500 dark:text-text-400 break-words">
+                Your Gmail credentials have been successfully uploaded and
+                authenticated.
+              </p>
+            </div>
+          </div>
+          <Button
+            variant="destructive"
+            type="button"
+            onClick={async () => {
+              handleRevokeAccess(
+                connectorExists,
+                setPopup,
+                existingCredential,
+                refreshCredentials
+              );
+            }}
+          >
+            Revoke Access
+          </Button>
+        </div>
+      </div>
+    );
+  }
+
+  // If no credentials are uploaded, show message to complete step 1 first
+  if (
+    !localServiceAccountData?.service_account_email &&
+    !localAppCredentialData?.client_id
+  ) {
+    return (
+      <div>
+        <SectionHeader>Gmail Authentication</SectionHeader>
+        <div className="mt-4">
+          <div className="flex items-start py-3 px-4 bg-yellow-50/30 dark:bg-yellow-900/5 rounded">
+            <FiAlertTriangle className="text-yellow-500 h-5 w-5 mr-2 mt-0.5 flex-shrink-0" />
+            <p className="text-sm">
+              Please complete Step 1 by uploading either OAuth credentials or a
+              Service Account key before proceeding with authentication.
+            </p>
+          </div>
+        </div>
+      </div>
     );
   }
 
   if (localServiceAccountData?.service_account_email) {
     return (
       <div>
-        <Formik
-          initialValues={{
-            google_primary_admin: user?.email || "",
-          }}
-          validationSchema={Yup.object().shape({
-            google_primary_admin: Yup.string()
-              .email("Must be a valid email")
-              .required("Required"),
-          })}
-          onSubmit={async (values, formikHelpers) => {
-            formikHelpers.setSubmitting(true);
-            try {
-              const response = await fetch(
-                "/api/manage/admin/connector/gmail/service-account-credential",
-                {
-                  method: "PUT",
-                  headers: {
-                    "Content-Type": "application/json",
-                  },
-                  body: JSON.stringify({
-                    google_primary_admin: values.google_primary_admin,
-                  }),
-                }
-              );
+        <div className="mt-4">
+          <Formik
+            initialValues={{
+              google_primary_admin: user?.email || "",
+            }}
+            validationSchema={Yup.object().shape({
+              google_primary_admin: Yup.string()
+                .email("Must be a valid email")
+                .required("Required"),
+            })}
+            onSubmit={async (values, formikHelpers) => {
+              formikHelpers.setSubmitting(true);
+              try {
+                const response = await fetch(
+                  "/api/manage/admin/connector/gmail/service-account-credential",
+                  {
+                    method: "PUT",
+                    headers: {
+                      "Content-Type": "application/json",
+                    },
+                    body: JSON.stringify({
+                      google_primary_admin: values.google_primary_admin,
+                    }),
+                  }
+                );
 
-              if (response.ok) {
+                if (response.ok) {
+                  setPopup({
+                    message: "Successfully created service account credential",
+                    type: "success",
+                  });
+                  refreshCredentials();
+                } else {
+                  const errorMsg = await response.text();
+                  setPopup({
+                    message: `Failed to create service account credential - ${errorMsg}`,
+                    type: "error",
+                  });
+                }
+              } catch (error) {
                 setPopup({
-                  message: "Successfully created service account credential",
-                  type: "success",
-                });
-                refreshCredentials();
-              } else {
-                const errorMsg = await response.text();
-                setPopup({
-                  message: `Failed to create service account credential - ${errorMsg}`,
+                  message: `Failed to create service account credential - ${error}`,
                   type: "error",
                 });
+              } finally {
+                formikHelpers.setSubmitting(false);
               }
-            } catch (error) {
-              setPopup({
-                message: `Failed to create service account credential - ${error}`,
-                type: "error",
-              });
-            } finally {
-              formikHelpers.setSubmitting(false);
-            }
-          }}
-        >
-          {({ isSubmitting }) => (
-            <Form>
-              <TextFormField
-                name="google_primary_admin"
-                label="Primary Admin Email:"
-                subtext="Enter the email of an admin/owner of the Google Organization that owns the Gmail account(s) you want to index."
-              />
-              <div className="flex">
-                <Button type="submit" disabled={isSubmitting}>
-                  Create Credential
-                </Button>
-              </div>
-            </Form>
-          )}
-        </Formik>
+            }}
+          >
+            {({ isSubmitting }) => (
+              <Form>
+                <TextFormField
+                  name="google_primary_admin"
+                  label="Primary Admin Email:"
+                  subtext="Enter the email of an admin/owner of the Google Organization that owns the Gmail account(s) you want to index."
+                />
+                <div className="flex">
+                  <Button type="submit" disabled={isSubmitting}>
+                    {isSubmitting ? "Creating..." : "Create Credential"}
+                  </Button>
+                </div>
+              </Form>
+            )}
+          </Formik>
+        </div>
       </div>
     );
   }
 
   if (localAppCredentialData?.client_id) {
     return (
-      <div className="text-sm mb-4">
-        <p className="mb-2">
-          Next, you must provide credentials via OAuth. This gives us read
-          access to the emails you have access to in your Gmail account.
-        </p>
+      <div>
+        <div className="bg-background-50/30 dark:bg-background-900/20 rounded mb-4">
+          <p className="text-sm">
+            Next, you need to authenticate with Gmail via OAuth. This gives us
+            read access to the emails you have access to in your Gmail account.
+          </p>
+        </div>
         <Button
+          disabled={isAuthenticating}
           onClick={async () => {
             setIsAuthenticating(true);
             try {
@@ -545,7 +654,6 @@ export const GmailAuthSection = ({
               setIsAuthenticating(false);
             }
           }}
-          disabled={isAuthenticating}
         >
           {isAuthenticating ? "Authenticating..." : "Authenticate with Gmail"}
         </Button>
@@ -553,11 +661,6 @@ export const GmailAuthSection = ({
     );
   }
 
-  // case where no keys have been uploaded in step 1
-  return (
-    <p className="text-sm">
-      Please upload either a OAuth Client Credential JSON or a Gmail Service
-      Account Key JSON in Step 1 before moving onto Step 2.
-    </p>
-  );
+  // This code path should not be reached with the new conditions above
+  return null;
 };

web/src/app/admin/connectors/[connector]/pages/gmail/GmailPage.tsx

@@ -173,6 +173,9 @@ export const GmailMain = () => {
         serviceAccountCredentialData={serviceAccountKeyData}
         isAdmin={isAdmin}
         onSuccess={handleRefresh}
+        existingAuthCredential={Boolean(
+          gmailPublicUploadedCredential || gmailServiceAccountCredential
+        )}
       />
 
       {isAdmin && hasUploadedCredentials && (

web/src/app/chat/ChatBanner.tsx

@@ -3,7 +3,7 @@
 import { SettingsContext } from "@/components/settings/SettingsProvider";
 import { useContext, useState, useRef, useLayoutEffect } from "react";
 import { ChevronDownIcon } from "@/components/icons/icons";
-import { MinimalMarkdown } from "@/components/chat/MinimalMarkdown";
+import MinimalMarkdown from "@/components/chat/MinimalMarkdown";
 
 export function ChatBanner() {
   const settings = useContext(SettingsContext);

web/src/app/chat/ChatPage.tsx

@@ -109,7 +109,6 @@ import {
 } from "@/components/resizable/constants";
 import FixedLogo from "../../components/logo/FixedLogo";
 
-import { MinimalMarkdown } from "@/components/chat/MinimalMarkdown";
 import ExceptionTraceModal from "@/components/modals/ExceptionTraceModal";
 
 import {
@@ -138,6 +137,7 @@ import { useSidebarShortcut } from "@/lib/browserUtilities";
 import { ConfirmEntityModal } from "@/components/modals/ConfirmEntityModal";
 import { ChatSearchModal } from "./chat_search/ChatSearchModal";
 import { ErrorBanner } from "./message/Resubmit";
+import MinimalMarkdown from "@/components/chat/MinimalMarkdown";
 
 const TEMP_USER_MESSAGE_ID = -1;
 const TEMP_ASSISTANT_MESSAGE_ID = -2;

web/src/app/chat/ChatPopup.tsx

@@ -6,6 +6,7 @@ import { Button } from "@/components/ui/button";
 import { useContext, useEffect, useState } from "react";
 import ReactMarkdown from "react-markdown";
 import remarkGfm from "remark-gfm";
+import { transformLinkUri } from "@/lib/utils";
 
 const ALL_USERS_INITIAL_POPUP_FLOW_COMPLETED =
   "allUsersInitialPopupFlowCompleted";
@@ -44,23 +45,26 @@ export function ChatPopup() {
   return (
     <Modal width="w-3/6 xl:w-[700px]" title={popupTitle}>
       <>
-        <ReactMarkdown
-          className="prose text-text-800 dark:text-neutral-100 max-w-full"
-          components={{
-            a: ({ node, ...props }) => (
-              <a
-                {...props}
-                className="text-link hover:text-link-hover"
-                target="_blank"
-                rel="noopener noreferrer"
-              />
-            ),
-            p: ({ node, ...props }) => <p {...props} className="text-sm" />,
-          }}
-          remarkPlugins={[remarkGfm]}
-        >
-          {popupContent}
-        </ReactMarkdown>
+        <div className="overflow-y-auto max-h-[90vh] py-8 px-4 text-left">
+          <ReactMarkdown
+            className="prose text-text-800 dark:text-neutral-100 max-w-full"
+            components={{
+              a: ({ node, ...props }) => (
+                <a
+                  {...props}
+                  className="text-link hover:text-link-hover"
+                  target="_blank"
+                  rel="noopener noreferrer"
+                />
+              ),
+              p: ({ node, ...props }) => <p {...props} className="text-sm" />,
+            }}
+            remarkPlugins={[remarkGfm]}
+            urlTransform={transformLinkUri}
+          >
+            {popupContent}
+          </ReactMarkdown>
+        </div>
 
         {showConsentError && (
           <p className="text-red-500 text-sm mt-2">

web/src/app/chat/message/AgenticMessage.tsx

@@ -53,6 +53,7 @@ import { copyAll, handleCopy } from "./copyingUtils";
 import { Button } from "@/components/ui/button";
 import { RefreshCw } from "lucide-react";
 import { ErrorBanner, Resubmit } from "./Resubmit";
+import { transformLinkUri } from "@/lib/utils";
 
 export const AgenticMessage = ({
   isStreamingQuestions,
@@ -336,6 +337,7 @@ export const AgenticMessage = ({
         }}
         remarkPlugins={[remarkGfm, remarkMath]}
         rehypePlugins={[[rehypePrism, { ignoreMissing: true }], rehypeKatex]}
+        urlTransform={transformLinkUri}
       >
         {finalAlternativeContent}
       </ReactMarkdown>
@@ -349,6 +351,7 @@ export const AgenticMessage = ({
         components={markdownComponents}
         remarkPlugins={[remarkGfm, remarkMath]}
         rehypePlugins={[[rehypePrism, { ignoreMissing: true }], rehypeKatex]}
+        urlTransform={transformLinkUri}
       >
         {streamedContent +
           (!isComplete && !secondLevelGenerating ? " [*]() " : "")}

web/src/app/chat/message/MemoizedTextComponents.tsx

@@ -160,8 +160,9 @@ export const MemoizedLink = memo(
 
     const handleMouseDown = () => {
       let url = href || rest.children?.toString();
-      if (url && !url.startsWith("http://") && !url.startsWith("https://")) {
-        // Try to construct a valid URL
+
+      if (url && !url.includes("://")) {
+        // Only add https:// if the URL doesn't already have a protocol
         const httpsUrl = `https://${url}`;
         try {
           new URL(httpsUrl);

web/src/app/chat/message/Messages.tsx

@@ -71,6 +71,7 @@ import remarkMath from "remark-math";
 import rehypeKatex from "rehype-katex";
 import "katex/dist/katex.min.css";
 import { copyAll, handleCopy } from "./copyingUtils";
+import { transformLinkUri } from "@/lib/utils";
 
 const TOOLS_WITH_CUSTOM_HANDLING = [
   SEARCH_TOOL_NAME,
@@ -348,7 +349,7 @@ export const AIMessage = ({
       a: anchorCallback,
       p: paragraphCallback,
       b: ({ node, className, children }: any) => {
-        return <span className={className}>||||{children}</span>;
+        return <span className={className}>{children}</span>;
       },
       code: ({ node, className, children }: any) => {
         const codeText = extractCodeText(
@@ -381,6 +382,7 @@ export const AIMessage = ({
         components={markdownComponents}
         remarkPlugins={[remarkGfm, remarkMath]}
         rehypePlugins={[[rehypePrism, { ignoreMissing: true }], rehypeKatex]}
+        urlTransform={transformLinkUri}
       >
         {finalContent}
       </ReactMarkdown>

web/src/app/chat/message/SubQuestionsDisplay.tsx

@@ -16,15 +16,15 @@ import ReactMarkdown from "react-markdown";
 import { MemoizedAnchor } from "./MemoizedTextComponents";
 import { MemoizedParagraph } from "./MemoizedTextComponents";
 import { extractCodeText, preprocessLaTeX } from "./codeUtils";
-
+import remarkGfm from "remark-gfm";
 import remarkMath from "remark-math";
 import rehypeKatex from "rehype-katex";
-import remarkGfm from "remark-gfm";
 import { CodeBlock } from "./CodeBlock";
 import { CheckIcon, ChevronDown } from "lucide-react";
 import { PHASE_MIN_MS, useStreamingMessages } from "./StreamingMessages";
 import { CirclingArrowIcon } from "@/components/icons/icons";
 import { handleCopy } from "./copyingUtils";
+import { transformLinkUri } from "@/lib/utils";
 
 export const StatusIndicator = ({ status }: { status: ToggleState }) => {
   return (
@@ -301,6 +301,7 @@ const SubQuestionDisplay: React.FC<{
         components={markdownComponents}
         remarkPlugins={[remarkGfm, remarkMath]}
         rehypePlugins={[rehypeKatex]}
+        urlTransform={transformLinkUri}
       >
         {finalContent}
       </ReactMarkdown>

web/src/app/chat/message/codeUtils.ts

@@ -62,19 +62,13 @@ export function extractCodeText(
 
 // We must preprocess LaTeX in the LLM output to avoid improper formatting
 export const preprocessLaTeX = (content: string) => {
-  // 1) Escape dollar signs used outside of LaTeX context
-  const escapedCurrencyContent = content.replace(
-    /\$(\d+(?:\.\d*)?)/g,
-    (_, p1) => `\\$${p1}`
-  );
-
-  // 2) Replace block-level LaTeX delimiters \[ \] with $$ $$
-  const blockProcessedContent = escapedCurrencyContent.replace(
+  // 1) Replace block-level LaTeX delimiters \[ \] with $$ $$
+  const blockProcessedContent = content.replace(
     /\\\[([\s\S]*?)\\\]/g,
     (_, equation) => `$$${equation}$$`
   );
 
-  // 3) Replace inline LaTeX delimiters \( \) with $ $
+  // 2) Replace inline LaTeX delimiters \( \) with $ $
   const inlineProcessedContent = blockProcessedContent.replace(
     /\\\(([\s\S]*?)\\\)/g,
     (_, equation) => `$${equation}$`
@@ -82,223 +76,3 @@ export const preprocessLaTeX = (content: string) => {
 
   return inlineProcessedContent;
 };
-
-interface MarkdownSegment {
-  type: "text" | "link" | "code" | "bold" | "italic" | "codeblock";
-  text: string; // The visible/plain text
-  raw: string; // The raw markdown including syntax
-  length: number; // Length of the visible text
-}
-
-export function parseMarkdownToSegments(markdown: string): MarkdownSegment[] {
-  if (!markdown) {
-    return [];
-  }
-
-  const segments: MarkdownSegment[] = [];
-  let currentIndex = 0;
-  const maxIterations = markdown.length * 2; // Prevent infinite loops
-  let iterations = 0;
-
-  while (currentIndex < markdown.length && iterations < maxIterations) {
-    iterations++;
-    let matched = false;
-
-    // Check for code blocks first (they take precedence)
-    const codeBlockMatch = markdown
-      .slice(currentIndex)
-      .match(/^```(\w*)\n([\s\S]*?)```/);
-    if (codeBlockMatch && codeBlockMatch[0]) {
-      const [fullMatch, , code] = codeBlockMatch;
-      segments.push({
-        type: "codeblock",
-        text: code || "",
-        raw: fullMatch,
-        length: (code || "").length,
-      });
-      currentIndex += fullMatch.length;
-      matched = true;
-      continue;
-    }
-
-    // Check for inline code
-    const inlineCodeMatch = markdown.slice(currentIndex).match(/^`([^`]+)`/);
-    if (inlineCodeMatch && inlineCodeMatch[0]) {
-      const [fullMatch, code] = inlineCodeMatch;
-      segments.push({
-        type: "code",
-        text: code || "",
-        raw: fullMatch,
-        length: (code || "").length,
-      });
-      currentIndex += fullMatch.length;
-      matched = true;
-      continue;
-    }
-
-    // Check for links
-    const linkMatch = markdown
-      .slice(currentIndex)
-      .match(/^\[([^\]]+)\]\(([^)]+)\)/);
-    if (linkMatch && linkMatch[0]) {
-      const [fullMatch, text] = linkMatch;
-      segments.push({
-        type: "link",
-        text: text || "",
-        raw: fullMatch,
-        length: (text || "").length,
-      });
-      currentIndex += fullMatch.length;
-      matched = true;
-      continue;
-    }
-
-    // Check for bold
-    const boldMatch = markdown
-      .slice(currentIndex)
-      .match(/^(\*\*|__)([^*_\n]*?)\1/);
-    if (boldMatch && boldMatch[0]) {
-      const [fullMatch, , text] = boldMatch;
-      segments.push({
-        type: "bold",
-        text: text || "",
-        raw: fullMatch,
-        length: (text || "").length,
-      });
-      currentIndex += fullMatch.length;
-      matched = true;
-      continue;
-    }
-
-    // Check for italic
-    const italicMatch = markdown
-      .slice(currentIndex)
-      .match(/^(\*|_)([^*_\n]+?)\1(?!\*|_)/);
-    if (italicMatch && italicMatch[0]) {
-      const [fullMatch, , text] = italicMatch;
-      segments.push({
-        type: "italic",
-        text: text || "",
-        raw: fullMatch,
-        length: (text || "").length,
-      });
-      currentIndex += fullMatch.length;
-      matched = true;
-      continue;
-    }
-
-    // If no matches were found, handle regular text
-    if (!matched) {
-      let nextSpecialChar = markdown.slice(currentIndex).search(/[`\[*_]/);
-      if (nextSpecialChar === -1) {
-        // No more special characters, add the rest as text
-        const text = markdown.slice(currentIndex);
-        if (text) {
-          segments.push({
-            type: "text",
-            text: text,
-            raw: text,
-            length: text.length,
-          });
-        }
-        break;
-      } else {
-        // Add the text up to the next special character
-        const text = markdown.slice(
-          currentIndex,
-          currentIndex + nextSpecialChar
-        );
-        if (text) {
-          segments.push({
-            type: "text",
-            text: text,
-            raw: text,
-            length: text.length,
-          });
-        }
-        currentIndex += nextSpecialChar;
-      }
-    }
-  }
-
-  return segments;
-}
-
-export function getMarkdownForSelection(
-  content: string,
-  selectedText: string
-): string {
-  const segments = parseMarkdownToSegments(content);
-
-  // Build plain text and create mapping to markdown segments
-  let plainText = "";
-  const markdownPieces: string[] = [];
-  let currentPlainIndex = 0;
-
-  segments.forEach((segment) => {
-    plainText += segment.text;
-    markdownPieces.push(segment.raw);
-    currentPlainIndex += segment.length;
-  });
-
-  // Find the selection in the plain text
-  const startIndex = plainText.indexOf(selectedText);
-  if (startIndex === -1) {
-    return selectedText;
-  }
-
-  const endIndex = startIndex + selectedText.length;
-
-  // Find which segments the selection spans
-  let currentIndex = 0;
-  let result = "";
-  let selectionStart = startIndex;
-  let selectionEnd = endIndex;
-
-  segments.forEach((segment) => {
-    const segmentStart = currentIndex;
-    const segmentEnd = segmentStart + segment.length;
-
-    // Check if this segment overlaps with the selection
-    if (segmentEnd > selectionStart && segmentStart < selectionEnd) {
-      // Calculate how much of this segment to include
-      const overlapStart = Math.max(0, selectionStart - segmentStart);
-      const overlapEnd = Math.min(segment.length, selectionEnd - segmentStart);
-
-      if (segment.type === "text") {
-        const textPortion = segment.text.slice(overlapStart, overlapEnd);
-        result += textPortion;
-      } else {
-        // For markdown elements, wrap just the selected portion with the appropriate markdown
-        const selectedPortion = segment.text.slice(overlapStart, overlapEnd);
-
-        switch (segment.type) {
-          case "bold":
-            result += `**${selectedPortion}**`;
-            break;
-          case "italic":
-            result += `*${selectedPortion}*`;
-            break;
-          case "code":
-            result += `\`${selectedPortion}\``;
-            break;
-          case "link":
-            // For links, we need to preserve the URL if it exists in the raw markdown
-            const urlMatch = segment.raw.match(/\]\((.*?)\)/);
-            const url = urlMatch ? urlMatch[1] : "";
-            result += `[${selectedPortion}](${url})`;
-            break;
-          case "codeblock":
-            result += `\`\`\`\n${selectedPortion}\n\`\`\``;
-            break;
-          default:
-            result += selectedPortion;
-        }
-      }
-    }
-
-    currentIndex += segment.length;
-  });
-
-  return result;
-}

web/src/components/admin/connectors/Field.tsx

@@ -33,6 +33,8 @@ import Link from "next/link";
 import { CheckboxField } from "@/components/ui/checkbox";
 import { CheckedState } from "@radix-ui/react-checkbox";
 
+import { transformLinkUri } from "@/lib/utils";
+
 export function SectionHeader({
   children,
 }: {
@@ -432,6 +434,7 @@ export const MarkdownFormField = ({
             <ReactMarkdown
               className="prose dark:prose-invert"
               remarkPlugins={[remarkGfm]}
+              urlTransform={transformLinkUri}
             >
               {field.value}
             </ReactMarkdown>

web/src/components/chat/MinimalMarkdown.tsx

@@ -4,19 +4,26 @@ import {
   MemoizedLink,
   MemoizedParagraph,
 } from "@/app/chat/message/MemoizedTextComponents";
-import React, { useMemo } from "react";
+import React, { useMemo, CSSProperties } from "react";
 import ReactMarkdown from "react-markdown";
 import remarkGfm from "remark-gfm";
+import rehypePrism from "rehype-prism-plus";
+import remarkMath from "remark-math";
+import rehypeKatex from "rehype-katex";
+import "katex/dist/katex.min.css";
+import { transformLinkUri } from "@/lib/utils";
 
 interface MinimalMarkdownProps {
   content: string;
   className?: string;
+  style?: CSSProperties;
 }
 
-export const MinimalMarkdown: React.FC<MinimalMarkdownProps> = ({
+export default function MinimalMarkdown({
   content,
   className = "",
-}) => {
+  style,
+}: MinimalMarkdownProps) {
   const markdownComponents = useMemo(
     () => ({
       a: MemoizedLink,
@@ -34,12 +41,16 @@ export const MinimalMarkdown: React.FC<MinimalMarkdownProps> = ({
   );
 
   return (
-    <ReactMarkdown
-      className={`w-full text-wrap break-word prose dark:prose-invert ${className}`}
-      components={markdownComponents}
-      remarkPlugins={[remarkGfm]}
-    >
-      {content}
-    </ReactMarkdown>
+    <div style={style || {}} className={`${className}`}>
+      <ReactMarkdown
+        className="prose dark:prose-invert max-w-full text-sm break-words"
+        components={markdownComponents}
+        rehypePlugins={[[rehypePrism, { ignoreMissing: true }], rehypeKatex]}
+        remarkPlugins={[remarkGfm, remarkMath]}
+        urlTransform={transformLinkUri}
+      >
+        {content}
+      </ReactMarkdown>
+    </div>
   );
-};
+}

web/src/components/chat/TextView.tsx

@@ -10,7 +10,7 @@ import {
 } from "@/components/ui/dialog";
 import { Download, XIcon, ZoomIn, ZoomOut } from "lucide-react";
 import { OnyxDocument } from "@/lib/search/interfaces";
-import { MinimalMarkdown } from "./MinimalMarkdown";
+import MinimalMarkdown from "@/components/chat/MinimalMarkdown";
 
 interface TextViewProps {
   presentingDocument: OnyxDocument;

web/src/components/context/ModalContext.tsx

@@ -4,6 +4,7 @@ import React, { createContext, useContext, useState, useCallback } from "react";
 import { NewTeamModal } from "../modals/NewTeamModal";
 import NewTenantModal from "../modals/NewTenantModal";
 import { User, NewTenantInfo } from "@/lib/types";
+import { NEXT_PUBLIC_CLOUD_ENABLED } from "@/lib/constants";
 
 type ModalContextType = {
   showNewTeamModal: boolean;
@@ -48,7 +49,7 @@ export const ModalProvider: React.FC<{
 
   // Render all application-wide modals
   const renderModals = () => {
-    if (!user) return null;
+    if (!user || !NEXT_PUBLIC_CLOUD_ENABLED) return null;
 
     return (
       <>

web/src/components/credentials/actions/CreateCredential.tsx

@@ -3,7 +3,10 @@ import { Button } from "@/components/ui/button";
 import { ValidSources } from "@/lib/types";
 import { FaAccusoft } from "react-icons/fa";
 import { submitCredential } from "@/components/admin/connectors/CredentialForm";
-import { TextFormField } from "@/components/admin/connectors/Field";
+import {
+  BooleanFormField,
+  TextFormField,
+} from "@/components/admin/connectors/Field";
 import { Form, Formik, FormikHelpers } from "formik";
 import { PopupSpec } from "@/components/admin/connectors/Popup";
 import { getSourceDocLink } from "@/lib/sources";
@@ -206,20 +209,31 @@ export default function CreateCredential({
               placeholder="(Optional) credential name.."
               label="Name:"
             />
-            {Object.entries(credentialTemplate).map(([key, val]) => (
-              <TextFormField
-                key={key}
-                name={key}
-                placeholder={val}
-                label={getDisplayNameForCredentialKey(key)}
-                type={
-                  key.toLowerCase().includes("token") ||
-                  key.toLowerCase().includes("password")
-                    ? "password"
-                    : "text"
-                }
-              />
-            ))}
+            {Object.entries(credentialTemplate).map(([key, val]) => {
+              if (typeof val === "boolean") {
+                return (
+                  <BooleanFormField
+                    key={key}
+                    name={key}
+                    label={getDisplayNameForCredentialKey(key)}
+                  />
+                );
+              }
+              return (
+                <TextFormField
+                  key={key}
+                  name={key}
+                  placeholder={val}
+                  label={getDisplayNameForCredentialKey(key)}
+                  type={
+                    key.toLowerCase().includes("token") ||
+                    key.toLowerCase().includes("password")
+                      ? "password"
+                      : "text"
+                  }
+                />
+              );
+            })}
             {!swapConnector && (
               <div className="mt-4 flex w-full flex-col sm:flex-row justify-between items-end">
                 <div className="w-full sm:w-3/4 mb-4 sm:mb-0">

web/src/components/search/results/AnswerSection.tsx

@@ -1,6 +1,6 @@
 import { Quote } from "@/lib/search/interfaces";
 import { ResponseSection, StatusOptions } from "./ResponseSection";
-import { MinimalMarkdown } from "@/components/chat/MinimalMarkdown";
+import MinimalMarkdown from "@/components/chat/MinimalMarkdown";
 
 const TEMP_STRING = "__$%^TEMP$%^__";
 

web/src/lib/connectors/credentials.ts

@@ -171,6 +171,7 @@ export interface SalesforceCredentialJson {
   sf_username: string;
   sf_password: string;
   sf_security_token: string;
+  is_sandbox: boolean;
 }
 
 export interface SharepointCredentialJson {
@@ -270,6 +271,7 @@ export const credentialTemplates: Record<ValidSources, any> = {
     sf_username: "",
     sf_password: "",
     sf_security_token: "",
+    is_sandbox: false,
   } as SalesforceCredentialJson,
   sharepoint: {
     sp_client_id: "",
@@ -452,6 +454,7 @@ export const credentialDisplayNames: Record<string, string> = {
   sf_username: "Salesforce Username",
   sf_password: "Salesforce Password",
   sf_security_token: "Salesforce Security Token",
+  is_sandbox: "Is Sandbox Environment",
 
   // Sharepoint
   sp_client_id: "SharePoint Client ID",

web/src/lib/constants.ts

@@ -91,3 +91,17 @@ export const NEXT_PUBLIC_INCLUDE_ERROR_POPUP_SUPPORT_LINK =
 
 export const NEXT_PUBLIC_STRIPE_PUBLISHABLE_KEY =
   process.env.NEXT_PUBLIC_STRIPE_PUBLISHABLE_KEY;
+
+// Add support for custom URL protocols in markdown links
+export const ALLOWED_URL_PROTOCOLS = [
+  "http:",
+  "https:",
+  "mailto:",
+  "tel:",
+  "slack:",
+  "vscode:",
+  "file:",
+  "sms:",
+  "spotify:",
+  "zoommtg:",
+];

web/src/lib/utils.ts

@@ -1,5 +1,6 @@
 import { clsx, type ClassValue } from "clsx";
 import { twMerge } from "tailwind-merge";
+import { ALLOWED_URL_PROTOCOLS } from "./constants";
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
@@ -8,3 +9,28 @@ export function cn(...inputs: ClassValue[]) {
 export const truncateString = (str: string, maxLength: number) => {
   return str.length > maxLength ? str.slice(0, maxLength - 1) + "..." : str;
 };
+
+/**
+ * Custom URL transformer function for ReactMarkdown
+ * Allows specific protocols to be used in markdown links
+ * We use this with the urlTransform prop in ReactMarkdown
+ */
+export function transformLinkUri(href: string) {
+  if (!href) return href;
+
+  const url = href.trim();
+  try {
+    const parsedUrl = new URL(url);
+    if (
+      ALLOWED_URL_PROTOCOLS.some((protocol) =>
+        parsedUrl.protocol.startsWith(protocol)
+      )
+    ) {
+      return url;
+    }
+  } catch (e) {
+    // If it's not a valid URL with protocol, return the original href
+    return href;
+  }
+  return href;
+}