feat: add auth login subcommand

codeflash/cli_cmds/cli.py

@@ -24,6 +24,8 @@ def parse_args() -> Namespace:
         args.no_pr = True
         args.worktree = True
         args.effort = "low"
+    if args.command == "auth":
+        return args
     return process_and_validate_cmd_args(args)
 
 
@@ -370,6 +372,11 @@ def _build_parser() -> ArgumentParser:
     subparsers.add_parser("vscode-install", help="Install the Codeflash VSCode extension")
     subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow")
 
+    auth_parser = subparsers.add_parser("auth", help="Authentication commands")
+    auth_subparsers = auth_parser.add_subparsers(dest="auth_command", help="Auth sub-commands")
+    auth_subparsers.add_parser("login", help="Log in to Codeflash via OAuth")
+    auth_subparsers.add_parser("status", help="Check authentication status")
+
     trace_optimize = subparsers.add_parser("optimize", help="Trace and optimize your project.")
 
     trace_optimize.add_argument(

codeflash/cli_cmds/cmd_auth.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import os
+
+import click
+
+from codeflash.cli_cmds.console import console
+from codeflash.cli_cmds.oauth_handler import perform_oauth_signin
+from codeflash.code_utils.env_utils import get_codeflash_api_key
+from codeflash.code_utils.shell_utils import save_api_key_to_rc
+from codeflash.either import is_successful
+
+
+def auth_login() -> None:
+    """Perform OAuth login and save the API key."""
+    try:
+        existing_api_key = get_codeflash_api_key()
+    except OSError:
+        existing_api_key = None
+
+    if existing_api_key:
+        display_key = f"{existing_api_key[:3]}****{existing_api_key[-4:]}"
+        console.print(f"[green]Already authenticated with API key {display_key}[/green]")
+        console.print("To re-authenticate, unset [bold]CODEFLASH_API_KEY[/bold] and run this command again.")
+        return
+
+    api_key = perform_oauth_signin()
+    if not api_key:
+        click.echo("Authentication failed.")
+        raise SystemExit(1)
+
+    result = save_api_key_to_rc(api_key)
+    if is_successful(result):
+        click.echo(result.unwrap())
+    else:
+        click.echo(result.failure())
+
+    os.environ["CODEFLASH_API_KEY"] = api_key
+    console.print("[green]Signed in successfully![/green]")
+
+
+def auth_status() -> None:
+    """Check and display current authentication status."""
+    try:
+        api_key = get_codeflash_api_key()
+    except OSError:
+        api_key = None
+
+    if api_key:
+        display_key = f"{api_key[:3]}****{api_key[-4:]}"
+        console.print(f"[green]Authenticated[/green] with API key {display_key}")
+    else:
+        console.print("[yellow]Not authenticated.[/yellow] Run [bold]codeflash auth login[/bold] to sign in.")

codeflash/main.py

@@ -34,7 +34,8 @@ def main() -> None:
     from codeflash.telemetry.sentry import init_sentry
 
     args = parse_args()
-    print_codeflash_banner()
+    if args.command != "auth":
+        print_codeflash_banner()
 
     # Check for newer version for all commands
     check_for_newer_minor_version()
@@ -47,7 +48,18 @@ def main() -> None:
         init_sentry(enabled=not disable_telemetry, exclude_errors=True)
         posthog_cf.initialize_posthog(enabled=not disable_telemetry)
 
-        if args.command == "init":
+        if args.command == "auth":
+            from codeflash.cli_cmds.cmd_auth import auth_login, auth_status
+
+            if args.auth_command == "login":
+                auth_login()
+            elif args.auth_command == "status":
+                auth_status()
+            else:
+                from codeflash.code_utils.code_utils import exit_with_message
+
+                exit_with_message("Usage: codeflash auth {login,status}", error_on_exit=True)
+        elif args.command == "init":
             from codeflash.cli_cmds.cmd_init import init_codeflash
 
             init_codeflash()

tests/test_cmd_auth.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from codeflash.cli_cmds.cmd_auth import auth_login
+from codeflash.either import Success
+
+
+class TestAuthLogin:
+    @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key")
+    @patch("codeflash.cli_cmds.cmd_auth.console")
+    def test_existing_api_key_skips_oauth(self, mock_console: MagicMock, mock_get_key: MagicMock) -> None:
+        mock_get_key.return_value = "cf-test1234abcd"
+
+        auth_login()
+
+        mock_console.print.assert_any_call("[green]Already authenticated with API key cf-****abcd[/green]")
+        mock_console.print.assert_any_call(
+            "To re-authenticate, unset [bold]CODEFLASH_API_KEY[/bold] and run this command again."
+        )
+
+    @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key")
+    @patch("codeflash.cli_cmds.cmd_auth.console")
+    def test_existing_api_key_oserror_treated_as_missing(
+        self, mock_console: MagicMock, mock_get_key: MagicMock
+    ) -> None:
+        mock_get_key.side_effect = OSError("permission denied")
+
+        with pytest.raises(SystemExit):
+            with patch("codeflash.cli_cmds.cmd_auth.perform_oauth_signin", return_value=None):
+                auth_login()
+
+    @patch("codeflash.cli_cmds.cmd_auth.perform_oauth_signin")
+    @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key", return_value="")
+    def test_oauth_failure_exits_with_code_1(self, mock_get_key: MagicMock, mock_oauth: MagicMock) -> None:
+        mock_oauth.return_value = None
+
+        with pytest.raises(SystemExit, match="1"):
+            auth_login()
+
+    @patch("codeflash.cli_cmds.cmd_auth.os")
+    @patch("codeflash.cli_cmds.cmd_auth.save_api_key_to_rc")
+    @patch("codeflash.cli_cmds.cmd_auth.perform_oauth_signin")
+    @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key", return_value="")
+    @patch("codeflash.cli_cmds.cmd_auth.console")
+    def test_successful_oauth_saves_key(
+        self,
+        mock_console: MagicMock,
+        mock_get_key: MagicMock,
+        mock_oauth: MagicMock,
+        mock_save: MagicMock,
+        mock_os: MagicMock,
+    ) -> None:
+        mock_oauth.return_value = "cf-newkey12345678"
+        mock_save.return_value = Success("API key saved to ~/.zshrc")
+
+        auth_login()
+
+        mock_save.assert_called_once_with("cf-newkey12345678")
+        mock_os.environ.__setitem__.assert_called_once_with("CODEFLASH_API_KEY", "cf-newkey12345678")
+        mock_console.print.assert_called_with("[green]Signed in successfully![/green]")
+
+    @patch("codeflash.cli_cmds.cmd_auth.os")
+    @patch("codeflash.cli_cmds.cmd_auth.save_api_key_to_rc")
+    @patch("codeflash.cli_cmds.cmd_auth.perform_oauth_signin")
+    @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key", return_value="")
+    @patch("codeflash.cli_cmds.cmd_auth.console")
+    def test_windows_oauth_saves_key(
+        self,
+        mock_console: MagicMock,
+        mock_get_key: MagicMock,
+        mock_oauth: MagicMock,
+        mock_save: MagicMock,
+        mock_os: MagicMock,
+    ) -> None:
+        mock_oauth.return_value = "cf-newkey12345678"
+        mock_os.name = "nt"
+        mock_save.return_value = Success("API key saved")
+
+        auth_login()
+
+        mock_save.assert_called_once_with("cf-newkey12345678")
+        mock_os.environ.__setitem__.assert_called_once_with("CODEFLASH_API_KEY", "cf-newkey12345678")
+
+
+class TestAuthSubcommandParsing:
+    def test_auth_login_parses(self) -> None:
+        from codeflash.cli_cmds.cli import _build_parser
+
+        _build_parser.cache_clear()
+        parser = _build_parser()
+        args = parser.parse_args(["auth", "login"])
+        assert args.command == "auth"
+        assert args.auth_command == "login"
+
+    def test_auth_without_subcommand(self) -> None:
+        from codeflash.cli_cmds.cli import _build_parser
+
+        _build_parser.cache_clear()
+        parser = _build_parser()
+        args = parser.parse_args(["auth"])
+        assert args.command == "auth"
+        assert args.auth_command is None