Merge pull request #1065 from airweave-ai/fix/token-manager-init

Fix: Create token manager based on oauth type not auth config class

Author: marc-rutzou

backend/airweave/platform/sync/factory.py

@@ -17,7 +17,6 @@
 from airweave.core.logging import ContextualLogger, LoggerConfigurator, logger
 from airweave.core.sync_cursor_service import sync_cursor_service
 from airweave.platform.auth_providers._base import BaseAuthProvider
-from airweave.platform.auth_providers.auth_result import AuthProviderMode
 from airweave.platform.destinations._base import BaseDestination
 from airweave.platform.entities._base import BaseEntity
 from airweave.platform.locator import resource_locator
@@ -473,34 +472,26 @@ async def _setup_token_manager(
     ) -> None:
         """Set up token manager for OAuth sources."""
         short_name = source_connection_data["short_name"]
-        auth_config_class_name = source_connection_data.get("auth_config_class")
         source_model = source_connection_data.get("source_model")
 
-        # Determine if we should create a token manager
+        # Determine if we should create a token manager based on oauth_type
         should_create_token_manager = False
 
-        # Case 1: Sources with OAuth2AuthConfig or its subclasses
-        if auth_config_class_name:
-            try:
-                # Get the auth config class
-                auth_config_class = resource_locator.get_auth_config(auth_config_class_name)
-
-                # Check if it's a subclass of OAuth2AuthConfig
-                from airweave.platform.configs.auth import OAuth2AuthConfig
-
-                if issubclass(auth_config_class, OAuth2AuthConfig):
-                    should_create_token_manager = True
-            except Exception as e:
-                logger.warning(f"Could not check auth config class for {short_name}: {str(e)}")
+        if source_model and hasattr(source_model, "oauth_type") and source_model.oauth_type:
+            # Import OAuthType enum
+            from airweave.schemas.source_connection import OAuthType
 
-        # Case 2: OAuth sources without auth_config_class (e.g., Asana, Google Calendar)
-        # These sources still need token management for refresh
-        elif source_model and hasattr(source_model, "oauth_type") and source_model.oauth_type:
-            # Check if we have OAuth credentials (dict with access_token)
-            if isinstance(source_credentials, dict) and "access_token" in source_credentials:
+            # Only create token manager for sources that support token refresh
+            if source_model.oauth_type in (OAuthType.WITH_REFRESH, OAuthType.WITH_ROTATING_REFRESH):
                 should_create_token_manager = True
                 logger.debug(
-                    f"✅ OAuth source {short_name} without auth_config_class will use token manager"
+                    f"✅ OAuth source {short_name} with oauth_type={source_model.oauth_type.value} "
+                    f"will use token manager for refresh"
+                )
+            else:
+                logger.debug(
+                    f"⏭️ Skipping token manager for {short_name} - "
+                    f"oauth_type={source_model.oauth_type} does not support token refresh"
                 )
 
         if should_create_token_manager:

backend/airweave/search/factory.py

@@ -890,17 +890,29 @@ async def _instantiate_federated_source(
                 ctx.logger.info(f"Proxy mode active for {source_connection.short_name}")
                 source_instance.set_http_client_factory(auth_config["http_client_factory"])
 
-            # Step 8: Setup token manager for OAuth sources
-            if source_model.oauth_type and isinstance(auth_config["credentials"], dict):
-                self._setup_token_manager(
-                    source_instance,
-                    db,
-                    source_connection,
-                    source_connection_data.get("integration_credential_id"),
-                    auth_config["credentials"],
-                    ctx,
-                    auth_provider_instance=auth_config.get("auth_provider_instance"),
-                )
+            # Step 8: Setup token manager for OAuth sources that support refresh
+            if source_model.oauth_type:
+                from airweave.schemas.source_connection import OAuthType
+
+                # Only create token manager for sources with refresh capability
+                if source_model.oauth_type in (
+                    OAuthType.WITH_REFRESH,
+                    OAuthType.WITH_ROTATING_REFRESH,
+                ):
+                    self._setup_token_manager(
+                        source_instance,
+                        db,
+                        source_connection,
+                        source_connection_data.get("integration_credential_id"),
+                        auth_config["credentials"],
+                        ctx,
+                        auth_provider_instance=auth_config.get("auth_provider_instance"),
+                    )
+                else:
+                    ctx.logger.debug(
+                        f"⏭️ Skipping token manager for {source_connection.short_name} - "
+                        f"oauth_type={source_model.oauth_type} does not support token refresh"
+                    )
 
             ctx.logger.info(
                 f"Successfully instantiated federated source: {source_connection.short_name}"