feat: fetch and store more data about okta authorization server (#3894)

Author: heitorado

lib/crewai/src/crewai/cli/authentication/main.py

@@ -1,5 +1,5 @@
 import time
-from typing import Any
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 import webbrowser
 
 from pydantic import BaseModel, Field
@@ -13,6 +13,8 @@
 
 console = Console()
 
+TOauth2Settings = TypeVar("TOauth2Settings", bound="Oauth2Settings")
+
 
 class Oauth2Settings(BaseModel):
     provider: str = Field(
@@ -28,22 +30,36 @@ class Oauth2Settings(BaseModel):
         description="OAuth2 audience value, typically used to identify the target API or resource.",
         default=None,
     )
+    extra: dict[str, Any] = Field(
+        description="Extra configuration for the OAuth2 provider.",
+        default={},
+    )
 
     @classmethod
-    def from_settings(cls):
+    def from_settings(cls: type[TOauth2Settings]) -> TOauth2Settings:
+        """Create an Oauth2Settings instance from the CLI settings."""
+
         settings = Settings()
 
         return cls(
             provider=settings.oauth2_provider,
             domain=settings.oauth2_domain,
             client_id=settings.oauth2_client_id,
             audience=settings.oauth2_audience,
+            extra=settings.oauth2_extra,
         )
 
 
+if TYPE_CHECKING:
+    from crewai.cli.authentication.providers.base_provider import BaseProvider
+
+
 class ProviderFactory:
     @classmethod
-    def from_settings(cls, settings: Oauth2Settings | None = None):
+    def from_settings(
+        cls: type["ProviderFactory"],  # noqa: UP037
+        settings: Oauth2Settings | None = None,
+    ) -> "BaseProvider":  # noqa: UP037
         settings = settings or Oauth2Settings.from_settings()
 
         import importlib
@@ -53,11 +69,11 @@ def from_settings(cls, settings: Oauth2Settings | None = None):
         )
         provider = getattr(module, f"{settings.provider.capitalize()}Provider")
 
-        return provider(settings)
+        return cast("BaseProvider", provider(settings))
 
 
 class AuthenticationCommand:
-    def __init__(self):
+    def __init__(self) -> None:
         self.token_manager = TokenManager()
         self.oauth2_provider = ProviderFactory.from_settings()
 
@@ -84,7 +100,7 @@ def _get_device_code(self) -> dict[str, Any]:
             timeout=20,
         )
         response.raise_for_status()
-        return response.json()
+        return cast(dict[str, Any], response.json())
 
     def _display_auth_instructions(self, device_code_data: dict[str, str]) -> None:
         """Display the authentication instructions to the user."""

lib/crewai/src/crewai/cli/authentication/providers/base_provider.py

@@ -24,3 +24,7 @@ def get_audience(self) -> str: ...
 
     @abstractmethod
     def get_client_id(self) -> str: ...
+
+    def get_required_fields(self) -> list[str]:
+        """Returns which provider-specific fields inside the "extra" dict will be required"""
+        return []

lib/crewai/src/crewai/cli/authentication/providers/okta.py

@@ -3,16 +3,16 @@
 
 class OktaProvider(BaseProvider):
     def get_authorize_url(self) -> str:
-        return f"https://{self.settings.domain}/oauth2/default/v1/device/authorize"
+        return f"{self._oauth2_base_url()}/v1/device/authorize"
 
     def get_token_url(self) -> str:
-        return f"https://{self.settings.domain}/oauth2/default/v1/token"
+        return f"{self._oauth2_base_url()}/v1/token"
 
     def get_jwks_url(self) -> str:
-        return f"https://{self.settings.domain}/oauth2/default/v1/keys"
+        return f"{self._oauth2_base_url()}/v1/keys"
 
     def get_issuer(self) -> str:
-        return f"https://{self.settings.domain}/oauth2/default"
+        return self._oauth2_base_url().removesuffix("/oauth2")
 
     def get_audience(self) -> str:
         if self.settings.audience is None:
@@ -27,3 +27,16 @@ def get_client_id(self) -> str:
                 "Client ID is required. Please set it in the configuration."
             )
         return self.settings.client_id
+
+    def get_required_fields(self) -> list[str]:
+        return ["authorization_server_name", "using_org_auth_server"]
+
+    def _oauth2_base_url(self) -> str:
+        using_org_auth_server = self.settings.extra.get("using_org_auth_server", False)
+
+        if using_org_auth_server:
+            base_url = f"https://{self.settings.domain}/oauth2"
+        else:
+            base_url = f"https://{self.settings.domain}/oauth2/{self.settings.extra.get('authorization_server_name', 'default')}"
+
+        return f"{base_url}"

lib/crewai/src/crewai/cli/command.py

@@ -11,18 +11,18 @@
 
 
 class BaseCommand:
-    def __init__(self):
+    def __init__(self) -> None:
         self._telemetry = Telemetry()
         self._telemetry.set_tracer()
 
 
 class PlusAPIMixin:
-    def __init__(self, telemetry):
+    def __init__(self, telemetry: Telemetry) -> None:
         try:
             telemetry.set_tracer()
             self.plus_api_client = PlusAPI(api_key=get_auth_token())
         except Exception:
-            self._deploy_signup_error_span = telemetry.deploy_signup_error_span()
+            telemetry.deploy_signup_error_span()
             console.print(
                 "Please sign up/login to CrewAI+ before using the CLI.",
                 style="bold red",

lib/crewai/src/crewai/cli/config.py

@@ -2,6 +2,7 @@
 from logging import getLogger
 from pathlib import Path
 import tempfile
+from typing import Any
 
 from pydantic import BaseModel, Field
 
@@ -136,7 +137,12 @@ class Settings(BaseModel):
         default=DEFAULT_CLI_SETTINGS["oauth2_domain"],
     )
 
-    def __init__(self, config_path: Path | None = None, **data):
+    oauth2_extra: dict[str, Any] = Field(
+        description="Extra configuration for the OAuth2 provider.",
+        default={},
+    )
+
+    def __init__(self, config_path: Path | None = None, **data: dict[str, Any]) -> None:
         """Load Settings from config path with fallback support"""
         if config_path is None:
             config_path = get_writable_config_path()

lib/crewai/src/crewai/cli/enterprise/main.py

@@ -1,9 +1,10 @@
-from typing import Any
+from typing import Any, cast
 
 import requests
 from requests.exceptions import JSONDecodeError, RequestException
 from rich.console import Console
 
+from crewai.cli.authentication.main import Oauth2Settings, ProviderFactory
 from crewai.cli.command import BaseCommand
 from crewai.cli.settings.main import SettingsCommand
 from crewai.cli.version import get_crewai_version
@@ -13,7 +14,7 @@
 
 
 class EnterpriseConfigureCommand(BaseCommand):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.settings_command = SettingsCommand()
 
@@ -54,25 +55,12 @@ def _fetch_oauth_config(self, enterprise_url: str) -> dict[str, Any]:
             except JSONDecodeError as e:
                 raise ValueError(f"Invalid JSON response from {oauth_endpoint}") from e
 
-            required_fields = [
-                "audience",
-                "domain",
-                "device_authorization_client_id",
-                "provider",
-            ]
-            missing_fields = [
-                field for field in required_fields if field not in oauth_config
-            ]
-
-            if missing_fields:
-                raise ValueError(
-                    f"Missing required fields in OAuth2 configuration: {', '.join(missing_fields)}"
-                )
+            self._validate_oauth_config(oauth_config)
 
             console.print(
                 "✅ Successfully retrieved OAuth2 configuration", style="green"
             )
-            return oauth_config
+            return cast(dict[str, Any], oauth_config)
 
         except RequestException as e:
             raise ValueError(f"Failed to connect to enterprise URL: {e!s}") from e
@@ -89,6 +77,7 @@ def _update_oauth_settings(
                 "oauth2_audience": oauth_config["audience"],
                 "oauth2_client_id": oauth_config["device_authorization_client_id"],
                 "oauth2_domain": oauth_config["domain"],
+                "oauth2_extra": oauth_config["extra"],
             }
 
             console.print("🔄 Updating local OAuth2 configuration...")
@@ -99,3 +88,38 @@ def _update_oauth_settings(
 
         except Exception as e:
             raise ValueError(f"Failed to update OAuth2 settings: {e!s}") from e
+
+    def _validate_oauth_config(self, oauth_config: dict[str, Any]) -> None:
+        required_fields = [
+            "audience",
+            "domain",
+            "device_authorization_client_id",
+            "provider",
+            "extra",
+        ]
+
+        missing_basic_fields = [
+            field for field in required_fields if field not in oauth_config
+        ]
+        missing_provider_specific_fields = [
+            field
+            for field in self._get_provider_specific_fields(oauth_config["provider"])
+            if field not in oauth_config.get("extra", {})
+        ]
+
+        if missing_basic_fields:
+            raise ValueError(
+                f"Missing required fields in OAuth2 configuration: [{', '.join(missing_basic_fields)}]"
+            )
+
+        if missing_provider_specific_fields:
+            raise ValueError(
+                f"Missing authentication provider required fields in OAuth2 configuration: [{', '.join(missing_provider_specific_fields)}] (Configured provider: '{oauth_config['provider']}')"
+            )
+
+    def _get_provider_specific_fields(self, provider_name: str) -> list[str]:
+        provider = ProviderFactory.from_settings(
+            Oauth2Settings(provider=provider_name, client_id="dummy", domain="dummy")
+        )
+
+        return provider.get_required_fields()

lib/crewai/src/crewai/cli/git.py

@@ -3,7 +3,7 @@
 
 
 class Repository:
-    def __init__(self, path="."):
+    def __init__(self, path: str = ".") -> None:
         self.path = path
 
         if not self.is_git_installed():

lib/crewai/src/crewai/cli/plus_api.py

@@ -1,3 +1,4 @@
+from typing import Any
 from urllib.parse import urljoin
 
 import requests
@@ -36,19 +37,21 @@ def __init__(self, api_key: str) -> None:
             str(settings.enterprise_base_url) or DEFAULT_CREWAI_ENTERPRISE_URL
         )
 
-    def _make_request(self, method: str, endpoint: str, **kwargs) -> requests.Response:
+    def _make_request(
+        self, method: str, endpoint: str, **kwargs: Any
+    ) -> requests.Response:
         url = urljoin(self.base_url, endpoint)
         session = requests.Session()
         session.trust_env = False
         return session.request(method, url, headers=self.headers, **kwargs)
 
-    def login_to_tool_repository(self):
+    def login_to_tool_repository(self) -> requests.Response:
         return self._make_request("POST", f"{self.TOOLS_RESOURCE}/login")
 
-    def get_tool(self, handle: str):
+    def get_tool(self, handle: str) -> requests.Response:
         return self._make_request("GET", f"{self.TOOLS_RESOURCE}/{handle}")
 
-    def get_agent(self, handle: str):
+    def get_agent(self, handle: str) -> requests.Response:
         return self._make_request("GET", f"{self.AGENTS_RESOURCE}/{handle}")
 
     def publish_tool(
@@ -58,8 +61,8 @@ def publish_tool(
         version: str,
         description: str | None,
         encoded_file: str,
-        available_exports: list[str] | None = None,
-    ):
+        available_exports: list[dict[str, Any]] | None = None,
+    ) -> requests.Response:
         params = {
             "handle": handle,
             "public": is_public,
@@ -111,28 +114,32 @@ def delete_crew_by_uuid(self, uuid: str) -> requests.Response:
     def list_crews(self) -> requests.Response:
         return self._make_request("GET", self.CREWS_RESOURCE)
 
-    def create_crew(self, payload) -> requests.Response:
+    def create_crew(self, payload: dict[str, Any]) -> requests.Response:
         return self._make_request("POST", self.CREWS_RESOURCE, json=payload)
 
     def get_organizations(self) -> requests.Response:
         return self._make_request("GET", self.ORGANIZATIONS_RESOURCE)
 
-    def initialize_trace_batch(self, payload) -> requests.Response:
+    def initialize_trace_batch(self, payload: dict[str, Any]) -> requests.Response:
         return self._make_request(
             "POST",
             f"{self.TRACING_RESOURCE}/batches",
             json=payload,
             timeout=30,
         )
 
-    def initialize_ephemeral_trace_batch(self, payload) -> requests.Response:
+    def initialize_ephemeral_trace_batch(
+        self, payload: dict[str, Any]
+    ) -> requests.Response:
         return self._make_request(
             "POST",
             f"{self.EPHEMERAL_TRACING_RESOURCE}/batches",
             json=payload,
         )
 
-    def send_trace_events(self, trace_batch_id: str, payload) -> requests.Response:
+    def send_trace_events(
+        self, trace_batch_id: str, payload: dict[str, Any]
+    ) -> requests.Response:
         return self._make_request(
             "POST",
             f"{self.TRACING_RESOURCE}/batches/{trace_batch_id}/events",
@@ -141,7 +148,7 @@ def send_trace_events(self, trace_batch_id: str, payload) -> requests.Response:
         )
 
     def send_ephemeral_trace_events(
-        self, trace_batch_id: str, payload
+        self, trace_batch_id: str, payload: dict[str, Any]
     ) -> requests.Response:
         return self._make_request(
             "POST",
@@ -150,7 +157,9 @@ def send_ephemeral_trace_events(
             timeout=30,
         )
 
-    def finalize_trace_batch(self, trace_batch_id: str, payload) -> requests.Response:
+    def finalize_trace_batch(
+        self, trace_batch_id: str, payload: dict[str, Any]
+    ) -> requests.Response:
         return self._make_request(
             "PATCH",
             f"{self.TRACING_RESOURCE}/batches/{trace_batch_id}/finalize",
@@ -159,7 +168,7 @@ def finalize_trace_batch(self, trace_batch_id: str, payload) -> requests.Respons
         )
 
     def finalize_ephemeral_trace_batch(
-        self, trace_batch_id: str, payload
+        self, trace_batch_id: str, payload: dict[str, Any]
     ) -> requests.Response:
         return self._make_request(
             "PATCH",

lib/crewai/src/crewai/cli/settings/main.py

@@ -34,7 +34,7 @@ def list(self) -> None:
             current_value = getattr(self.settings, field_name)
             description = field_info.description or "No description available"
             display_value = (
-                str(current_value) if current_value is not None else "Not set"
+                str(current_value) if current_value not in [None, {}] else "Not set"
             )
 
             table.add_row(field_name, display_value, description)