feat(platform): add organization parameter to OAuth authorization redirect

specifications/SPEC_PLATFORM_SERVICE.md

@@ -409,14 +409,15 @@ The Platform module provides foundational services but does not directly expose
 
 ### 6.1 Configuration Parameters
 
-| Parameter | Type | Default | Description | Required |
-| ------------------------------- | ---- | --------------------------------- | --------------------------- | -------- |
-| `api_root` | str | `https://platform.aignostics.com` | Base URL of Aignostics API | Yes |
-| `audience` | str | Environment-specific | OAuth audience claim | Yes |
-| `scope` | str | `offline_access` | OAuth scopes required | Yes |
-| `cache_dir` | str | User cache directory | Directory for token storage | No |
-| `request_timeout_seconds` | int | 30 | API request timeout | No |
-| `authorization_backoff_seconds` | int | 3 | Retry backoff time | No |
+| Parameter | Type | Default | Description | Required |
+| ------------------------------- | ---- | --------------------------------- | ---------------------------- | -------- |
+| `api_root` | str | `https://platform.aignostics.com` | Base URL of Aignostics API | Yes |
+| `audience` | str | Environment-specific | OAuth audience claim | Yes |
+| `scope` | str | `offline_access` | OAuth scopes required | Yes |
+| `organization_id` | str | None | Auth0 organization for OAuth | No |
+| `cache_dir` | str | User cache directory | Directory for token storage | No |
+| `request_timeout_seconds` | int | 30 | API request timeout | No |
+| `authorization_backoff_seconds` | int | 3 | Retry backoff time | No |
 
 ### 6.2 Environment Variables
 
@@ -432,6 +433,7 @@ The Platform module provides foundational services but does not directly expose
 | `AIGNOSTICS_DEVICE_URL` | Custom device authorization URL | `https://custom.auth0.com/oauth/device/code` |
 | `AIGNOSTICS_JWS_JSON_URL` | Custom JWS key set URL | `https://custom.auth0.com/.well-known/jwks.json` |
 | `AIGNOSTICS_CLIENT_ID_INTERACTIVE` | Interactive flow client ID | `interactive_client_123` |
+| `AIGNOSTICS_ORGANIZATION_ID` | Auth0 organization for OAuth | `my-organization` |
 | `AIGNOSTICS_REFRESH_TOKEN` | Long-lived refresh token | `refresh_token_value` |
 | `AIGNOSTICS_CACHE_DIR` | Custom cache directory | `/custom/cache/path` |
 | `AIGNOSTICS_REQUEST_TIMEOUT_SECONDS` | API request timeout | `60` |
@@ -492,7 +494,7 @@ The Platform module provides foundational services but does not directly expose
 
 ### 9.1 Key Algorithms and Business Logic
 
-- **PKCE Flow**: OAuth 2.0 Authorization Code flow with Proof Key for Code Exchange for enhanced security in public clients
+- **PKCE Flow**: OAuth 2.0 Authorization Code flow with Proof Key for Code Exchange for enhanced security in public clients. Supports optional `AIGNOSTICS_ORGANIZATION_ID` parameter for Auth0 organization-specific authentication flows
 - **Token Caching**: File-based token persistence with expiration tracking and automatic cleanup
 - **Health Monitoring**: Multi-layer health checks including public endpoint availability and authenticated API access
 

src/aignostics/platform/_authentication.py

@@ -333,10 +333,18 @@ def _perform_authorization_code_with_pkce_flow() -> str: # noqa: C901
  redirect_uri=settings().redirect_uri,
  pkce="S256",
  )
+ auth_params = {
+ "access_type": "offline",
+ "audience": settings().audience,
+ }
+ organization_id = settings().organization_id
+
+ if organization_id:
+ auth_params["organization"] = organization_id
+
  authorization_url, _ = session.authorization_url(
  settings().authorization_base_url,
- access_type="offline",
- audience=settings().audience,
+ **auth_params,
  )
 
  authentication_result = AuthenticationResult()

src/aignostics/platform/_settings.py

@@ -206,6 +206,10 @@ def profile_edit_url(self) -> str:
  ]
  client_id_interactive: Annotated[str, Field(description="OAuth client ID for interactive flows")]
 
+ organization_id: Annotated[
+ str | None, Field(description="Optional Auth0 organization ID parameter for the /authorize OAuth endpoint")
+ ] = None
+
  @computed_field # type: ignore[prop-decorator]
  @property
  def tenant_domain(self) -> str:

tests/aignostics/platform/authentication_test.py

@@ -61,6 +61,7 @@
  settings.auth_retry_attempts = 3
  settings.auth_jwk_set_cache_ttl = 300
  settings.refresh_token = None
+ settings.organization_id = None
  mock_settings.return_value = settings
  yield mock_settings
 
@@ -339,25 +340,22 @@
 class TestAuthorizationCodeFlow:
  """Test cases for the authorization code flow with PKCE."""
 
- @pytest.mark.unit
  @staticmethod
- def test_perform_authorization_code_flow_success(record_property, mock_settings) -> None:
- """Test successful authorization code flow with PKCE."""
- record_property("tested-item-id", "SPEC-PLATFORM-SERVICE")
- # Mock OAuth session
+ def _run_pkce_flow() -> tuple[str | None, MagicMock]:
+ """Set up common PKCE flow mocks, run the flow, and return (token, mock_session).
+
+ Returns:
+ tuple: The returned token and the OAuth2Session mock for further assertions.
+ """
  mock_session = MagicMock(spec=OAuth2Session)
  mock_session.authorization_url.return_value = ("https://test.auth/authorize?code_challenge=abc", None)
  mock_session.fetch_token.return_value = {"access_token": "pkce.token"}
 
- # Mock HTTP server
  mock_server = MagicMock()
-
- # Setup mocks for the redirect URI parsing
  mock_redirect_parsed = MagicMock()
  mock_redirect_parsed.hostname = "localhost"
  mock_redirect_parsed.port = 8000
 
- # Create a custom HTTPServer mock implementation that simulates a callback
  class MockHTTPServer:
  def __init__(self, *args, **kwargs) -> None:
  pass
@@ -365,10 +363,9 @@
  def __enter__(self) -> MagicMock:
  return mock_server
 
  def __exit__(self, *args) -> None:
  pass
 
- # Create a mock for the auth result
  mock_auth_result = MagicMock()
  mock_auth_result.token = "pkce.token" # noqa: S105 - Test credential
  mock_auth_result.error = None
@@ -379,20 +376,19 @@
  patch("urllib.parse.urlparse", return_value=mock_redirect_parsed),
  patch("aignostics.platform._authentication.AuthenticationResult", return_value=mock_auth_result),
  ):
- # Simulate a successful server response by making handle_request set the token
- def handle_request_side_effect():
- # This simulates what the HTTP handler would do on success
- mock_auth_result.token = "pkce.token" # noqa: S105 - Test credential
-
- mock_server.handle_request.side_effect = handle_request_side_effect
-
- # Call the function under test
+ mock_server.handle_request.side_effect = lambda: None
  token = _perform_authorization_code_with_pkce_flow()
 
- # Assertions
- assert token == "pkce.token" # noqa: S105 - Test credential
- mock_server.handle_request.assert_called_once()
- mock_session.authorization_url.assert_called_once()
+ return token, mock_session
+
+ @pytest.mark.unit
+ @staticmethod
+ def test_perform_authorization_code_flow_success(record_property, mock_settings) -> None:
+ """Test successful authorization code flow with PKCE."""
+ record_property("tested-item-id", "SPEC-PLATFORM-SERVICE")
+ token, mock_session = TestAuthorizationCodeFlow._run_pkce_flow()
+ assert token == "pkce.token" # noqa: S105 - Test credential
+ mock_session.authorization_url.assert_called_once()
 
  @pytest.mark.unit
  @staticmethod
@@ -466,6 +462,27 @@
  with pytest.raises(RuntimeError, match=AUTHENTICATION_FAILED):
  _perform_authorization_code_with_pkce_flow()
 
+ @pytest.mark.unit
+ @staticmethod
+ def test_perform_authorization_code_flow_with_organization(record_property, mock_settings) -> None:
+ """Test authorization code flow includes organization parameter when set."""
+ record_property("tested-item-id", "SPEC-PLATFORM-SERVICE")
+ mock_settings.return_value.organization_id = "test-org"
+ token, mock_session = TestAuthorizationCodeFlow._run_pkce_flow()
+ assert token == "pkce.token" # noqa: S105
+ assert mock_session.authorization_url.call_args[1].get("organization") == "test-org"
+
+ @pytest.mark.unit
+ @staticmethod
+ def test_perform_authorization_code_flow_without_organization(record_property, mock_settings) -> None:
+ """Test authorization code flow omits organization parameter when unset."""
+ record_property("tested-item-id", "SPEC-PLATFORM-SERVICE")
+ # organization_id is None by default in mock_settings fixture
+ assert mock_settings.return_value.organization_id is None
+ token, mock_session = TestAuthorizationCodeFlow._run_pkce_flow()
+ assert token == "pkce.token" # noqa: S105
+ assert "organization" not in mock_session.authorization_url.call_args[1]
+
 
 class TestDeviceFlow:
  """Test cases for the device flow authentication."""