feat: add multiple provider support (#3089)

* Remove `crewai signup` command, update docs

* Add `Settings.clear()` and clear settings before each login

* Add pyjwt

* Remove print statement from ToolCommand.login()

* Remove auth0 dependency

* Update docs

tests/cli/authentication/test_auth_main.py

@@ -1,103 +1,419 @@
-import unittest
-from unittest.mock import MagicMock, patch
-
+import pytest
+from datetime import datetime, timedelta
 import requests
-
+from unittest.mock import MagicMock, patch, call
 from crewai.cli.authentication.main import AuthenticationCommand
+from crewai.cli.authentication.constants import (
+    AUTH0_AUDIENCE,
+    AUTH0_CLIENT_ID,
+    AUTH0_DOMAIN,
+    WORKOS_DOMAIN,
+    WORKOS_CLI_CONNECT_APP_ID,
+    WORKOS_ENVIRONMENT_ID,
+)
 
 
-class TestAuthenticationCommand(unittest.TestCase):
-    def setUp(self):
+class TestAuthenticationCommand:
+    def setup_method(self):
         self.auth_command = AuthenticationCommand()
 
-    @patch("crewai.cli.authentication.main.requests.post")
-    def test_get_device_code(self, mock_post):
-        mock_response = MagicMock()
-        mock_response.json.return_value = {
-            "device_code": "123456",
-            "user_code": "ABCDEF",
-            "verification_uri_complete": "https://example.com",
-            "interval": 5,
-        }
-        mock_post.return_value = mock_response
-
-        device_code_data = self.auth_command._get_device_code()
-
-        self.assertEqual(device_code_data["device_code"], "123456")
-        self.assertEqual(device_code_data["user_code"], "ABCDEF")
-        self.assertEqual(
-            device_code_data["verification_uri_complete"], "https://example.com"
-        )
-        self.assertEqual(device_code_data["interval"], 5)
-
+    @pytest.mark.parametrize(
+        "user_provider,expected_urls",
+        [
+            (
+                "auth0",
+                {
+                    "device_code_url": f"https://{AUTH0_DOMAIN}/oauth/device/code",
+                    "token_url": f"https://{AUTH0_DOMAIN}/oauth/token",
+                    "client_id": AUTH0_CLIENT_ID,
+                    "audience": AUTH0_AUDIENCE,
+                },
+            ),
+            (
+                "workos",
+                {
+                    "device_code_url": f"https://{WORKOS_DOMAIN}/oauth2/device_authorization",
+                    "token_url": f"https://{WORKOS_DOMAIN}/oauth2/token",
+                    "client_id": WORKOS_CLI_CONNECT_APP_ID,
+                },
+            ),
+        ],
+    )
+    @patch(
+        "crewai.cli.authentication.main.AuthenticationCommand._determine_user_provider"
+    )
+    @patch("crewai.cli.authentication.main.AuthenticationCommand._get_device_code")
+    @patch(
+        "crewai.cli.authentication.main.AuthenticationCommand._display_auth_instructions"
+    )
+    @patch("crewai.cli.authentication.main.AuthenticationCommand._poll_for_token")
     @patch("crewai.cli.authentication.main.console.print")
-    @patch("crewai.cli.authentication.main.webbrowser.open")
-    def test_display_auth_instructions(self, mock_open, mock_print):
+    def test_login(
+        self,
+        mock_console_print,
+        mock_poll,
+        mock_display,
+        mock_get_device,
+        mock_determine_provider,
+        user_provider,
+        expected_urls,
+    ):
+        mock_determine_provider.return_value = user_provider
+        mock_get_device.return_value = {
+            "device_code": "test_code",
+            "user_code": "123456",
+        }
+
+        self.auth_command.login()
+
+        mock_console_print.assert_called_once_with(
+            "Signing in to CrewAI Enterprise...\n", style="bold blue"
+        )
+        mock_determine_provider.assert_called_once()
+        mock_get_device.assert_called_once_with(
+            expected_urls["client_id"],
+            expected_urls["device_code_url"],
+            expected_urls.get("audience", None),
+        )
+        mock_display.assert_called_once_with(
+            {"device_code": "test_code", "user_code": "123456"}
+        )
+        mock_poll.assert_called_once_with(
+            {"device_code": "test_code", "user_code": "123456"},
+            expected_urls["client_id"],
+            expected_urls["token_url"],
+        )
+
+    @patch("crewai.cli.authentication.main.webbrowser")
+    @patch("crewai.cli.authentication.main.console.print")
+    def test_display_auth_instructions(self, mock_console_print, mock_webbrowser):
         device_code_data = {
-            "verification_uri_complete": "https://example.com",
-            "user_code": "ABCDEF",
+            "verification_uri_complete": "https://example.com/auth",
+            "user_code": "123456",
         }
 
         self.auth_command._display_auth_instructions(device_code_data)
 
-        mock_print.assert_any_call("1. Navigate to: ", "https://example.com")
-        mock_print.assert_any_call("2. Enter the following code: ", "ABCDEF")
-        mock_open.assert_called_once_with("https://example.com")
+        expected_calls = [
+            call("1. Navigate to: ", "https://example.com/auth"),
+            call("2. Enter the following code: ", "123456"),
+        ]
+        mock_console_print.assert_has_calls(expected_calls)
+        mock_webbrowser.open.assert_called_once_with("https://example.com/auth")
+
+    @pytest.mark.parametrize(
+        "user_provider,jwt_config",
+        [
+            (
+                "auth0",
+                {
+                    "jwks_url": f"https://{AUTH0_DOMAIN}/.well-known/jwks.json",
+                    "issuer": f"https://{AUTH0_DOMAIN}/",
+                    "audience": AUTH0_AUDIENCE,
+                },
+            ),
+            (
+                "workos",
+                {
+                    "jwks_url": f"https://{WORKOS_DOMAIN}/oauth2/jwks",
+                    "issuer": f"https://{WORKOS_DOMAIN}",
+                    "audience": WORKOS_ENVIRONMENT_ID,
+                },
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("has_expiration", [True, False])
+    @patch("crewai.cli.authentication.main.validate_jwt_token")
+    @patch("crewai.cli.authentication.main.TokenManager.save_tokens")
+    def test_validate_and_save_token(
+        self,
+        mock_save_tokens,
+        mock_validate_jwt,
+        user_provider,
+        jwt_config,
+        has_expiration,
+    ):
+        self.auth_command.user_provider = user_provider
+        token_data = {"access_token": "test_access_token", "id_token": "test_id_token"}
+
+        if has_expiration:
+            future_timestamp = int((datetime.now() + timedelta(days=100)).timestamp())
+            decoded_token = {"exp": future_timestamp}
+        else:
+            decoded_token = {}
+
+        mock_validate_jwt.return_value = decoded_token
+
+        self.auth_command._validate_and_save_token(token_data)
+
+        mock_validate_jwt.assert_called_once_with(
+            jwt_token="test_access_token",
+            jwks_url=jwt_config["jwks_url"],
+            issuer=jwt_config["issuer"],
+            audience=jwt_config["audience"],
+        )
+
+        if has_expiration:
+            mock_save_tokens.assert_called_once_with(
+                "test_access_token", future_timestamp
+            )
+        else:
+            mock_save_tokens.assert_called_once_with("test_access_token", 0)
 
     @patch("crewai.cli.tools.main.ToolCommand")
-    @patch("crewai.cli.authentication.main.requests.post")
-    @patch("crewai.cli.authentication.main.validate_token")
+    @patch("crewai.cli.authentication.main.Settings")
     @patch("crewai.cli.authentication.main.console.print")
-    def test_poll_for_token_success(
-        self, mock_print, mock_validate_token, mock_post, mock_tool
+    def test_login_to_tool_repository_success(
+        self, mock_console_print, mock_settings, mock_tool_command
     ):
+        mock_tool_instance = MagicMock()
+        mock_tool_command.return_value = mock_tool_instance
+
+        mock_settings_instance = MagicMock()
+        mock_settings_instance.org_name = "Test Org"
+        mock_settings_instance.org_uuid = "test-uuid-123"
+        mock_settings.return_value = mock_settings_instance
+
+        self.auth_command._login_to_tool_repository()
+
+        mock_tool_command.assert_called_once()
+        mock_tool_instance.login.assert_called_once()
+
+        expected_calls = [
+            call(
+                "Now logging you in to the Tool Repository... ",
+                style="bold blue",
+                end="",
+            ),
+            call("Success!\n", style="bold green"),
+            call(
+                "You are authenticated to the tool repository as [bold cyan]'Test Org'[/bold cyan] (test-uuid-123)",
+                style="green",
+            ),
+        ]
+        mock_console_print.assert_has_calls(expected_calls)
+
+    @patch("crewai.cli.tools.main.ToolCommand")
+    @patch("crewai.cli.authentication.main.console.print")
+    def test_login_to_tool_repository_error(
+        self, mock_console_print, mock_tool_command
+    ):
+        mock_tool_instance = MagicMock()
+        mock_tool_instance.login.side_effect = Exception("Tool repository error")
+        mock_tool_command.return_value = mock_tool_instance
+
+        self.auth_command._login_to_tool_repository()
+
+        mock_tool_command.assert_called_once()
+        mock_tool_instance.login.assert_called_once()
+
+        expected_calls = [
+            call(
+                "Now logging you in to the Tool Repository... ",
+                style="bold blue",
+                end="",
+            ),
+            call(
+                "\n[bold yellow]Warning:[/bold yellow] Authentication with the Tool Repository failed.",
+                style="yellow",
+            ),
+            call(
+                "Other features will work normally, but you may experience limitations with downloading and publishing tools.\nRun [bold]crewai login[/bold] to try logging in again.\n",
+                style="yellow",
+            ),
+        ]
+        mock_console_print.assert_has_calls(expected_calls)
+
+    @pytest.mark.parametrize(
+        "api_response,expected_provider",
+        [
+            ({"provider": "auth0"}, "auth0"),
+            ({"provider": "workos"}, "workos"),
+            ({"provider": "none"}, "workos"),  # Default to workos for any other value
+            (
+                {},
+                "workos",
+            ),  # Default to workos if no provider key is sent in the response
+        ],
+    )
+    @patch("crewai.cli.authentication.main.PlusAPI")
+    @patch("crewai.cli.authentication.main.console.print")
+    @patch("builtins.input", return_value="test@example.com")
+    def test_determine_user_provider_success(
+        self,
+        mock_input,
+        mock_console_print,
+        mock_plus_api,
+        api_response,
+        expected_provider,
+    ):
+        mock_api_instance = MagicMock()
         mock_response = MagicMock()
         mock_response.status_code = 200
-        mock_response.json.return_value = {
-            "id_token": "TOKEN",
-            "access_token": "ACCESS_TOKEN",
-        }
-        mock_post.return_value = mock_response
+        mock_response.json.return_value = api_response
+        mock_api_instance._make_request.return_value = mock_response
+        mock_plus_api.return_value = mock_api_instance
 
-        mock_instance = mock_tool.return_value
-        mock_instance.login.return_value = None
+        result = self.auth_command._determine_user_provider()
 
-        self.auth_command._poll_for_token({"device_code": "123456"})
+        mock_input.assert_called_once()
 
-        mock_validate_token.assert_called_once_with("TOKEN")
-        mock_print.assert_called_once_with(
-            "\n[bold green]Welcome to CrewAI Enterprise![/bold green]\n"
+        mock_plus_api.assert_called_once_with("")
+        mock_api_instance._make_request.assert_called_once_with(
+            "GET", "/crewai_plus/api/v1/me/provider?email=test%40example.com"
         )
 
-    @patch("crewai.cli.authentication.main.requests.post")
+        assert result == expected_provider
+
+    @patch("crewai.cli.authentication.main.PlusAPI")
     @patch("crewai.cli.authentication.main.console.print")
-    def test_poll_for_token_error(self, mock_print, mock_post):
+    @patch("builtins.input", return_value="test@example.com")
+    def test_determine_user_provider_error(
+        self, mock_input, mock_console_print, mock_plus_api
+    ):
+        mock_api_instance = MagicMock()
+        mock_response = MagicMock()
+        mock_response.status_code = 500
+        mock_api_instance._make_request.return_value = mock_response
+        mock_plus_api.return_value = mock_api_instance
+
+        with pytest.raises(SystemExit):
+            self.auth_command._determine_user_provider()
+
+        mock_input.assert_called_once()
+
+        mock_plus_api.assert_called_once_with("")
+        mock_api_instance._make_request.assert_called_once_with(
+            "GET", "/crewai_plus/api/v1/me/provider?email=test%40example.com"
+        )
+
+        mock_console_print.assert_has_calls(
+            [
+                call(
+                    "Enter your CrewAI Enterprise account email: ",
+                    style="bold blue",
+                    end="",
+                ),
+                call(
+                    "Error: Failed to authenticate with crewai enterprise. Ensure that you are using the latest crewai version and please try again. If the problem persists, contact support@crewai.com.",
+                    style="red",
+                ),
+            ]
+        )
+
+    @patch("requests.post")
+    def test_get_device_code(self, mock_post):
         mock_response = MagicMock()
-        mock_response.status_code = 400
         mock_response.json.return_value = {
-            "error": "invalid_request",
-            "error_description": "Invalid request",
+            "device_code": "test_device_code",
+            "user_code": "123456",
+            "verification_uri_complete": "https://example.com/auth",
         }
         mock_post.return_value = mock_response
 
-        with self.assertRaises(requests.HTTPError):
-            self.auth_command._poll_for_token({"device_code": "123456"})
+        result = self.auth_command._get_device_code(
+            client_id="test_client",
+            device_code_url="https://example.com/device",
+            audience="test_audience",
+        )
 
-        mock_print.assert_not_called()
+        mock_post.assert_called_once_with(
+            url="https://example.com/device",
+            data={
+                "client_id": "test_client",
+                "scope": "openid",
+                "audience": "test_audience",
+            },
+            timeout=20,
+        )
 
-    @patch("crewai.cli.authentication.main.requests.post")
-    @patch("crewai.cli.authentication.main.console.print")
-    def test_poll_for_token_timeout(self, mock_print, mock_post):
-        mock_response = MagicMock()
-        mock_response.status_code = 400
-        mock_response.json.return_value = {
-            "error": "authorization_pending",
-            "error_description": "Authorization pending",
+        assert result == {
+            "device_code": "test_device_code",
+            "user_code": "123456",
+            "verification_uri_complete": "https://example.com/auth",
         }
-        mock_post.return_value = mock_response
 
-        self.auth_command._poll_for_token({"device_code": "123456", "interval": 0.01})
+    @patch("requests.post")
+    @patch("crewai.cli.authentication.main.console.print")
+    def test_poll_for_token_success(self, mock_console_print, mock_post):
+        mock_response_success = MagicMock()
+        mock_response_success.status_code = 200
+        mock_response_success.json.return_value = {
+            "access_token": "test_access_token",
+            "id_token": "test_id_token",
+        }
+        mock_post.return_value = mock_response_success
 
-        mock_print.assert_called_once_with(
+        device_code_data = {"device_code": "test_device_code", "interval": 1}
+
+        with (
+            patch.object(
+                self.auth_command, "_validate_and_save_token"
+            ) as mock_validate,
+            patch.object(
+                self.auth_command, "_login_to_tool_repository"
+            ) as mock_tool_login,
+        ):
+            self.auth_command._poll_for_token(
+                device_code_data, "test_client", "https://example.com/token"
+            )
+
+            mock_post.assert_called_once_with(
+                "https://example.com/token",
+                data={
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                    "device_code": "test_device_code",
+                    "client_id": "test_client",
+                },
+                timeout=30,
+            )
+
+            mock_validate.assert_called_once()
+            mock_tool_login.assert_called_once()
+
+            expected_calls = [
+                call("\nWaiting for authentication... ", style="bold blue", end=""),
+                call("Success!", style="bold green"),
+                call("\n[bold green]Welcome to CrewAI Enterprise![/bold green]\n"),
+            ]
+            mock_console_print.assert_has_calls(expected_calls)
+
+    @patch("requests.post")
+    @patch("crewai.cli.authentication.main.console.print")
+    def test_poll_for_token_timeout(self, mock_console_print, mock_post):
+        mock_response_pending = MagicMock()
+        mock_response_pending.status_code = 400
+        mock_response_pending.json.return_value = {"error": "authorization_pending"}
+        mock_post.return_value = mock_response_pending
+
+        device_code_data = {
+            "device_code": "test_device_code",
+            "interval": 0.1,  # Short interval for testing
+        }
+
+        self.auth_command._poll_for_token(
+            device_code_data, "test_client", "https://example.com/token"
+        )
+
+        mock_console_print.assert_any_call(
             "Timeout: Failed to get the token. Please try again.", style="bold red"
         )
+
+    @patch("requests.post")
+    def test_poll_for_token_error(self, mock_post):
+        """Test the method to poll for token (error path)."""
+        # Setup mock to return error
+        mock_response_error = MagicMock()
+        mock_response_error.status_code = 400
+        mock_response_error.json.return_value = {
+            "error": "access_denied",
+            "error_description": "User denied access",
+        }
+        mock_post.return_value = mock_response_error
+
+        device_code_data = {"device_code": "test_device_code", "interval": 1}
+
+        with pytest.raises(requests.HTTPError):
+            self.auth_command._poll_for_token(
+                device_code_data, "test_client", "https://example.com/token"
+            )

tests/cli/authentication/test_utils.py

@@ -1,31 +1,110 @@
 import json
+import jwt
 import unittest
 from datetime import datetime, timedelta
 from unittest.mock import MagicMock, patch
 
 from cryptography.fernet import Fernet
 
-from crewai.cli.authentication.utils import TokenManager, validate_token
+from crewai.cli.authentication.utils import TokenManager, validate_jwt_token
 
 
+@patch("crewai.cli.authentication.utils.PyJWKClient", return_value=MagicMock())
+@patch("crewai.cli.authentication.utils.jwt")
 class TestValidateToken(unittest.TestCase):
-    @patch("crewai.cli.authentication.utils.AsymmetricSignatureVerifier")
-    @patch("crewai.cli.authentication.utils.TokenVerifier")
-    def test_validate_token(self, mock_token_verifier, mock_asymmetric_verifier):
-        mock_verifier_instance = mock_token_verifier.return_value
-        mock_id_token = "mock_id_token"
+    def test_validate_jwt_token(self, mock_jwt, mock_pyjwkclient):
+        mock_jwt.decode.return_value = {"exp": 1719859200}
 
-        validate_token(mock_id_token)
+        # Create signing key object mock with a .key attribute
+        mock_pyjwkclient.return_value.get_signing_key_from_jwt.return_value = MagicMock(
+            key="mock_signing_key"
+        )
 
-        mock_asymmetric_verifier.assert_called_once_with(
-            "https://crewai.us.auth0.com/.well-known/jwks.json"
+        decoded_token = validate_jwt_token(
+            jwt_token="aaaaa.bbbbbb.cccccc",
+            jwks_url="https://mock_jwks_url",
+            issuer="https://mock_issuer",
+            audience="app_id_xxxx",
         )
-        mock_token_verifier.assert_called_once_with(
-            signature_verifier=mock_asymmetric_verifier.return_value,
-            issuer="https://crewai.us.auth0.com/",
-            audience="DEVC5Fw6NlRoSzmDCcOhVq85EfLBjKa8",
+
+        mock_jwt.decode.assert_called_once_with(
+            "aaaaa.bbbbbb.cccccc",
+            "mock_signing_key",
+            algorithms=["RS256"],
+            audience="app_id_xxxx",
+            issuer="https://mock_issuer",
+            options={
+                "verify_signature": True,
+                "verify_exp": True,
+                "verify_nbf": True,
+                "verify_iat": True,
+                "require": ["exp", "iat", "iss", "aud", "sub"],
+            },
         )
-        mock_verifier_instance.verify.assert_called_once_with(mock_id_token)
+        mock_pyjwkclient.assert_called_once_with("https://mock_jwks_url")
+        self.assertEqual(decoded_token, {"exp": 1719859200})
+
+    def test_validate_jwt_token_expired(self, mock_jwt, mock_pyjwkclient):
+        mock_jwt.decode.side_effect = jwt.ExpiredSignatureError
+        with self.assertRaises(Exception):
+            validate_jwt_token(
+                jwt_token="aaaaa.bbbbbb.cccccc",
+                jwks_url="https://mock_jwks_url",
+                issuer="https://mock_issuer",
+                audience="app_id_xxxx",
+            )
+
+    def test_validate_jwt_token_invalid_audience(self, mock_jwt, mock_pyjwkclient):
+        mock_jwt.decode.side_effect = jwt.InvalidAudienceError
+        with self.assertRaises(Exception):
+            validate_jwt_token(
+                jwt_token="aaaaa.bbbbbb.cccccc",
+                jwks_url="https://mock_jwks_url",
+                issuer="https://mock_issuer",
+                audience="app_id_xxxx",
+            )
+
+    def test_validate_jwt_token_invalid_issuer(self, mock_jwt, mock_pyjwkclient):
+        mock_jwt.decode.side_effect = jwt.InvalidIssuerError
+        with self.assertRaises(Exception):
+            validate_jwt_token(
+                jwt_token="aaaaa.bbbbbb.cccccc",
+                jwks_url="https://mock_jwks_url",
+                issuer="https://mock_issuer",
+                audience="app_id_xxxx",
+            )
+
+    def test_validate_jwt_token_missing_required_claims(
+        self, mock_jwt, mock_pyjwkclient
+    ):
+        mock_jwt.decode.side_effect = jwt.MissingRequiredClaimError
+        with self.assertRaises(Exception):
+            validate_jwt_token(
+                jwt_token="aaaaa.bbbbbb.cccccc",
+                jwks_url="https://mock_jwks_url",
+                issuer="https://mock_issuer",
+                audience="app_id_xxxx",
+            )
+
+    def test_validate_jwt_token_jwks_error(self, mock_jwt, mock_pyjwkclient):
+        mock_jwt.decode.side_effect = jwt.exceptions.PyJWKClientError
+        with self.assertRaises(Exception):
+            validate_jwt_token(
+                jwt_token="aaaaa.bbbbbb.cccccc",
+                jwks_url="https://mock_jwks_url",
+                issuer="https://mock_issuer",
+                audience="app_id_xxxx",
+            )
+
+    def test_validate_jwt_token_invalid_token(self, mock_jwt, mock_pyjwkclient):
+        mock_jwt.decode.side_effect = jwt.InvalidTokenError
+        with self.assertRaises(Exception):
+            validate_jwt_token(
+                jwt_token="aaaaa.bbbbbb.cccccc",
+                jwks_url="https://mock_jwks_url",
+                issuer="https://mock_issuer",
+                audience="app_id_xxxx",
+            )
 
 
 class TestTokenManager(unittest.TestCase):
@@ -62,9 +141,9 @@ class TestTokenManager(unittest.TestCase):
     @patch("crewai.cli.authentication.utils.TokenManager.save_secure_file")
     def test_save_tokens(self, mock_save):
         access_token = "test_token"
-        expires_in = 3600
+        expires_at = int((datetime.now() + timedelta(seconds=3600)).timestamp())
 
-        self.token_manager.save_tokens(access_token, expires_in)
+        self.token_manager.save_tokens(access_token, expires_at)
 
         mock_save.assert_called_once()
         args = mock_save.call_args[0]
@@ -73,11 +152,7 @@ class TestTokenManager(unittest.TestCase):
         data = json.loads(decrypted_data)
         self.assertEqual(data["access_token"], access_token)
         expiration = datetime.fromisoformat(data["expiration"])
-        self.assertAlmostEqual(
-            expiration,
-            datetime.now() + timedelta(seconds=expires_in),
-            delta=timedelta(seconds=1),
-        )
+        self.assertEqual(expiration, datetime.fromtimestamp(expires_at))
 
     @patch("crewai.cli.authentication.utils.TokenManager.read_secure_file")
     def test_get_token_valid(self, mock_read):

tests/cli/cli_test.py

@@ -13,7 +13,7 @@ from crewai.cli.cli import (
     deply_status,
     flow_add_crew,
     reset_memories,
-    signup,
+    login,
     test,
     train,
     version,
@@ -261,12 +261,12 @@ def test_test_invalid_string_iterations(evaluate_crew, runner):
 
 
 @mock.patch("crewai.cli.cli.AuthenticationCommand")
-def test_signup(command, runner):
+def test_login(command, runner):
     mock_auth = command.return_value
-    result = runner.invoke(signup)
+    result = runner.invoke(login)
 
     assert result.exit_code == 0
-    mock_auth.signup.assert_called_once()
+    mock_auth.login.assert_called_once()
 
 
 @mock.patch("crewai.cli.cli.DeployCommand")

tests/cli/tools/test_main.py

@@ -2,6 +2,7 @@ import os
 import tempfile
 import unittest
 import unittest.mock
+from datetime import datetime, timedelta
 from contextlib import contextmanager
 from unittest import mock
 from unittest.mock import MagicMock, patch
@@ -26,7 +27,9 @@ def in_temp_dir():
 
 @pytest.fixture
 def tool_command():
-    TokenManager().save_tokens("test-token", 36000)
+    TokenManager().save_tokens(
+        "test-token", (datetime.now() + timedelta(seconds=36000)).timestamp()
+    )
     tool_command = ToolCommand()
     with patch.object(tool_command, "login"):
         yield tool_command
@@ -57,7 +60,9 @@ def test_create_success(mock_subprocess, capsys, tool_command):
 @patch("crewai.cli.tools.main.subprocess.run")
 @patch("crewai.cli.plus_api.PlusAPI.get_tool")
 @patch("crewai.cli.tools.main.ToolCommand._print_current_organization")
-def test_install_success(mock_print_org, mock_get, mock_subprocess_run, capsys, tool_command):
+def test_install_success(
+    mock_print_org, mock_get, mock_subprocess_run, capsys, tool_command
+):
     mock_get_response = MagicMock()
     mock_get_response.status_code = 200
     mock_get_response.json.return_value = {
@@ -89,6 +94,7 @@ def test_install_success(mock_print_org, mock_get, mock_subprocess_run, capsys,
     # Verify _print_current_organization was called
     mock_print_org.assert_called_once()
 
+
 @patch("crewai.cli.tools.main.subprocess.run")
 @patch("crewai.cli.plus_api.PlusAPI.get_tool")
 def test_install_success_from_pypi(mock_get, mock_subprocess_run, capsys, tool_command):
@@ -169,7 +175,10 @@ def test_publish_when_not_in_sync(mock_is_synced, capsys, tool_command):
 )
 @patch("crewai.cli.plus_api.PlusAPI.publish_tool")
 @patch("crewai.cli.tools.main.git.Repository.is_synced", return_value=False)
-@patch("crewai.cli.tools.main.extract_available_exports", return_value=[{"name": "SampleTool"}])
+@patch(
+    "crewai.cli.tools.main.extract_available_exports",
+    return_value=[{"name": "SampleTool"}],
+)
 @patch("crewai.cli.tools.main.ToolCommand._print_current_organization")
 def test_publish_when_not_in_sync_and_force(
     mock_print_org,
@@ -223,7 +232,10 @@ def test_publish_when_not_in_sync_and_force(
 )
 @patch("crewai.cli.plus_api.PlusAPI.publish_tool")
 @patch("crewai.cli.tools.main.git.Repository.is_synced", return_value=True)
-@patch("crewai.cli.tools.main.extract_available_exports", return_value=[{"name": "SampleTool"}])
+@patch(
+    "crewai.cli.tools.main.extract_available_exports",
+    return_value=[{"name": "SampleTool"}],
+)
 def test_publish_success(
     mock_available_exports,
     mock_is_synced,
@@ -273,7 +285,10 @@ def test_publish_success(
     read_data=b"sample tarball content",
 )
 @patch("crewai.cli.plus_api.PlusAPI.publish_tool")
-@patch("crewai.cli.tools.main.extract_available_exports", return_value=[{"name": "SampleTool"}])
+@patch(
+    "crewai.cli.tools.main.extract_available_exports",
+    return_value=[{"name": "SampleTool"}],
+)
 def test_publish_failure(
     mock_available_exports,
     mock_publish,
@@ -311,7 +326,10 @@ def test_publish_failure(
     read_data=b"sample tarball content",
 )
 @patch("crewai.cli.plus_api.PlusAPI.publish_tool")
-@patch("crewai.cli.tools.main.extract_available_exports", return_value=[{"name": "SampleTool"}])
+@patch(
+    "crewai.cli.tools.main.extract_available_exports",
+    return_value=[{"name": "SampleTool"}],
+)
 def test_publish_api_error(
     mock_available_exports,
     mock_publish,
@@ -338,7 +356,6 @@ def test_publish_api_error(
     mock_publish.assert_called_once()
 
 
-
 @patch("crewai.cli.tools.main.Settings")
 def test_print_current_organization_with_org(mock_settings, capsys, tool_command):
     mock_settings_instance = MagicMock()