Merge PR sooperset#718: Feature/Discover the cloud-id from the provided OAuth token when in BYOAccessToken

Author: nszceta

src/mcp_atlassian/utils/oauth.py

@@ -143,7 +143,7 @@ def exchange_code_for_tokens(self, code: str) -> bool:
             self.expires_at = time.time() + token_data["expires_in"]
 
             # Get the cloud ID using the access token
-            self._get_cloud_id()
+            self.cloud_id = get_cloud_id(self.access_token)
 
             # Save the tokens
             self._save_tokens()
@@ -229,32 +229,6 @@ def ensure_valid_token(self) -> bool:
             return True
         return self.refresh_access_token()
 
-    def _get_cloud_id(self) -> None:
-        """Get the cloud ID for the Atlassian instance.
-
-        This method queries the accessible resources endpoint to get the cloud ID.
-        The cloud ID is needed for API calls with OAuth.
-        """
-        if not self.access_token:
-            logger.debug("No access token available to get cloud ID")
-            return
-
-        try:
-            headers = {"Authorization": f"Bearer {self.access_token}"}
-            response = requests.get(CLOUD_ID_URL, headers=headers)
-            response.raise_for_status()
-
-            resources = response.json()
-            if resources and len(resources) > 0:
-                # Use the first cloud site (most users have only one)
-                # For users with multiple sites, they might need to specify which one to use
-                self.cloud_id = resources[0]["id"]
-                logger.debug(f"Found cloud ID: {self.cloud_id}")
-            else:
-                logger.warning("No Atlassian sites found in the response")
-        except Exception as e:
-            logger.error(f"Failed to get cloud ID: {e}")
-
     def _get_keyring_username(self) -> str:
         """Get the keyring username for storing tokens.
 
@@ -465,15 +439,45 @@ def from_env(cls) -> Optional["BYOAccessTokenOAuthConfig"]:
             BYOAccessTokenOAuthConfig instance or None if required
             environment variables are missing.
         """
-        cloud_id = os.getenv("ATLASSIAN_OAUTH_CLOUD_ID")
+
         access_token = os.getenv("ATLASSIAN_OAUTH_ACCESS_TOKEN")
+        cloud_id = os.getenv("ATLASSIAN_OAUTH_CLOUD_ID") or get_cloud_id(access_token)
 
         if not all([cloud_id, access_token]):
             return None
 
         return cls(cloud_id=cloud_id, access_token=access_token)
 
 
+def get_cloud_id(access_token: str) -> str | None:
+    """Get the cloud ID for the Atlassian instance.
+
+    This method queries the accessible resources endpoint to get the cloud ID.
+    The cloud ID is needed for API calls with OAuth.
+    """
+    if not access_token:
+        logger.debug("No access token provided")
+        return None
+
+    try:
+        response = requests.get(
+            CLOUD_ID_URL, headers={"Authorization": f"Bearer {access_token}"}
+        )
+        response.raise_for_status()
+        resources = response.json() or []
+        cloud_id = resources[0].get("id") if resources else None
+
+        if cloud_id:
+            logger.debug(f"Found cloud ID: {cloud_id}")
+        else:
+            logger.warning("No Atlassian sites found in the response")
+
+        return cloud_id
+    except Exception as e:
+        logger.error(f"Failed to get cloud ID: {e}")
+        return None
+
+
 def get_oauth_config_from_env() -> OAuthConfig | BYOAccessTokenOAuthConfig | None:
     """Get the appropriate OAuth configuration from environment variables.
 

tests/unit/utils/test_oauth.py

@@ -139,7 +139,7 @@ def test_exchange_code_for_tokens_success(self, mock_post):
         mock_post.return_value = mock_response
 
         # Mock cloud ID retrieval and token saving
-        with patch.object(OAuthConfig, "_get_cloud_id") as mock_get_cloud_id:
+        with patch("mcp_atlassian.utils.oauth.get_cloud_id") as mock_get_cloud_id:
             with patch.object(OAuthConfig, "_save_tokens") as mock_save_tokens:
                 config = OAuthConfig(
                     client_id="test-client-id",
@@ -157,7 +157,7 @@ def test_exchange_code_for_tokens_success(self, mock_post):
 
                 # Verify calls
                 mock_post.assert_called_once()
-                mock_get_cloud_id.assert_called_once()
+                mock_get_cloud_id.assert_called_once_with("new-access-token")
                 mock_save_tokens.assert_called_once()
 
     @patch("requests.post")
@@ -284,42 +284,33 @@ def test_ensure_valid_token_needs_refresh_failure(self, mock_refresh):
 
     @patch("requests.get")
     def test_get_cloud_id_success(self, mock_get):
-        """Test _get_cloud_id success case."""
+        """Test get_cloud_id success case."""
         # Mock response
         mock_response = MagicMock()
         mock_response.status_code = 200
         mock_response.json.return_value = [{"id": "test-cloud-id", "name": "Test Site"}]
         mock_get.return_value = mock_response
 
-        config = OAuthConfig(
-            client_id="test-client-id",
-            client_secret="test-client-secret",
-            redirect_uri="https://example.com/callback",
-            scope="read:jira-work write:jira-work",
-            access_token="test-access-token",
-        )
-        config._get_cloud_id()
+        from mcp_atlassian.utils.oauth import get_cloud_id
+
+        cloud_id = get_cloud_id("test-access-token")
 
         # Check result
-        assert config.cloud_id == "test-cloud-id"
+        assert cloud_id == "test-cloud-id"
         mock_get.assert_called_once()
         headers = mock_get.call_args[1]["headers"]
         assert headers["Authorization"] == "Bearer test-access-token"
 
     @patch("requests.get")
     def test_get_cloud_id_no_access_token(self, mock_get):
-        """Test _get_cloud_id with no access token."""
-        config = OAuthConfig(
-            client_id="test-client-id",
-            client_secret="test-client-secret",
-            redirect_uri="https://example.com/callback",
-            scope="read:jira-work write:jira-work",
-        )
-        config._get_cloud_id()
+        """Test get_cloud_id with no access token."""
+        from mcp_atlassian.utils.oauth import get_cloud_id
+
+        cloud_id = get_cloud_id(None)
 
         # Should not make API call without token
         mock_get.assert_not_called()
-        assert config.cloud_id is None
+        assert cloud_id is None
 
     def test_get_keyring_username(self):
         """Test _get_keyring_username method."""