chore: clean up auth session (#452)

Author: caroljung-cg

src/codegen/cli/auth/auth_session.py

@@ -4,9 +4,10 @@
 import click
 import rich
 
-from codegen.cli.auth.auth_session import CodegenAuthenticatedSession
 from codegen.cli.auth.login import login_routine
-from codegen.cli.errors import AuthError, InvalidTokenError, NoTokenError
+from codegen.cli.auth.session import CodegenSession
+from codegen.cli.auth.token_manager import TokenManager, get_current_token
+from codegen.cli.errors import AuthError
 from codegen.cli.rich.pretty_print import pretty_print_error
 
 
@@ -15,22 +16,23 @@ def requires_auth(f: Callable) -> Callable:
 
     @functools.wraps(f)
     def wrapper(*args, **kwargs):
-        session = CodegenAuthenticatedSession.from_active_session()
+        session = CodegenSession.from_active_session()
 
         # Check for valid session
-        if not session.is_valid():
-            pretty_print_error(f"The session at path {session.repo_path} is missing or corrupt.\nPlease run 'codegen init' to re-initialize the project.")
+        if session is None or not session.is_valid():
+            pretty_print_error("There is currently no active session.\nPlease run 'codegen init' to initialize the project.")
             raise click.Abort()
 
-        try:
-            if not session.is_authenticated():
-                rich.print("[yellow]Not authenticated. Let's get you logged in first![/yellow]\n")
-                session = login_routine()
-        except (InvalidTokenError, NoTokenError) as e:
-            rich.print("[yellow]Authentication token is invalid or expired. Let's get you logged in again![/yellow]\n")
-            session = login_routine()
-        except AuthError as e:
-            raise click.ClickException(str(e))
+        if (token := get_current_token()) is None:
+            rich.print("[yellow]Not authenticated. Let's get you logged in first![/yellow]\n")
+            login_routine()
+        else:
+            try:
+                token_manager = TokenManager()
+                token_manager.authenticate_token(token)
+            except AuthError:
+                rich.print("[yellow]Authentication token is invalid or expired. Let's get you logged in again![/yellow]\n")
+                login_routine()
 
         return f(*args, session=session, **kwargs)
 

src/codegen/cli/auth/decorators.py

@@ -4,48 +4,45 @@
 import rich_click as click
 
 from codegen.cli.api.webapp_routes import USER_SECRETS_ROUTE
-from codegen.cli.auth.auth_session import CodegenAuthenticatedSession
 from codegen.cli.auth.token_manager import TokenManager
 from codegen.cli.env.global_env import global_env
 from codegen.cli.errors import AuthError
 
 
-def login_routine(token: str | None = None) -> CodegenAuthenticatedSession:
+def login_routine(token: str | None = None) -> str:
     """Guide user through login flow and return authenticated session.
 
     Args:
-        console: Optional console for output. Creates new one if not provided.
+        token: Codegen user access token associated with github account
 
     Returns:
-        CodegenSession: Authenticated session
+        str: The authenticated token
 
     Raises:
         click.ClickException: If login fails
 
     """
     # Try environment variable first
-
-    _token = token or global_env.CODEGEN_USER_ACCESS_TOKEN
+    token = token or global_env.CODEGEN_USER_ACCESS_TOKEN
 
     # If no token provided, guide user through browser flow
-    if not _token:
+    if not token:
         rich.print(f"Opening {USER_SECRETS_ROUTE} to get your authentication token...")
         webbrowser.open_new(USER_SECRETS_ROUTE)
-        _token = click.prompt("Please enter your authentication token from the browser", hide_input=False)
+        token = click.prompt("Please enter your authentication token from the browser", hide_input=False)
 
-    if not _token:
+    if not token:
         msg = "Token must be provided via CODEGEN_USER_ACCESS_TOKEN environment variable or manual input"
         raise click.ClickException(msg)
 
     # Validate and store token
-    token_manager = TokenManager()
-    session = CodegenAuthenticatedSession(token=_token)
-
     try:
-        session.assert_authenticated()
-        token_manager.save_token(_token)
+        token_manager = TokenManager()
+        token_manager.authenticate_token(token)
         rich.print(f"[green]✓ Stored token to:[/green] {token_manager.token_file}")
-        return session
+        rich.print("[cyan]📊 Hey![/cyan] We collect anonymous usage data to improve your experience 🔒")
+        rich.print("To opt out, set [green]telemetry_enabled = false[/green] in [cyan]~/.config/codegen-sh/analytics.json[/cyan] ✨")
+        return token
     except AuthError as e:
         msg = f"Error: {e!s}"
         raise click.ClickException(msg)

src/codegen/cli/auth/login.py

@@ -2,7 +2,9 @@
 import os
 from pathlib import Path
 
+from codegen.cli.api.client import RestAPI
 from codegen.cli.auth.constants import AUTH_FILE, CONFIG_DIR
+from codegen.cli.errors import AuthError
 
 
 class TokenManager:
@@ -19,6 +21,17 @@ def _ensure_config_dir(self):
         if not os.path.exists(self.config_dir):
             Path(self.config_dir).mkdir(parents=True, exist_ok=True)
 
+    def authenticate_token(self, token: str) -> None:
+        """Authenticate the token with the api."""
+        identity = RestAPI(token).identify()
+        if not identity:
+            msg = "No identity found for session"
+            raise AuthError(msg)
+        if identity.auth_context.status != "active":
+            msg = "Current session is not active. API Token may be invalid or may have expired."
+            raise AuthError(msg)
+        self.save_token(token)
+
     def save_token(self, token: str) -> None:
         """Save api token to disk."""
         try:

src/codegen/cli/auth/token_manager.py

@@ -5,14 +5,15 @@
 
 from codegen.cli.api.client import RestAPI
 from codegen.cli.auth.constants import PROMPTS_DIR
+from codegen.cli.auth.decorators import requires_auth
 from codegen.cli.auth.session import CodegenSession
+from codegen.cli.auth.token_manager import get_current_token
 from codegen.cli.codemod.convert import convert_to_cli
 from codegen.cli.errors import ServerError
 from codegen.cli.rich.codeblocks import format_command, format_path
 from codegen.cli.rich.pretty_print import pretty_print_error
 from codegen.cli.rich.spinners import create_spinner
 from codegen.cli.utils.default_code import DEFAULT_CODEMOD
-from codegen.cli.workspace.decorators import requires_init
 
 
 def get_prompts_dir() -> Path:
@@ -65,7 +66,7 @@ def make_relative(path: Path) -> str:
 
 
 @click.command(name="create")
-@requires_init
+@requires_auth
 @click.argument("name", type=str)
 @click.argument("path", type=click.Path(path_type=Path), default=Path.cwd())
 @click.option("--description", "-d", default=None, help="Description of what this codemod does.")
@@ -92,7 +93,7 @@ def create_command(session: CodegenSession, name: str, path: Path, description:
         if description:
             # Use API to generate implementation
             with create_spinner("Generating function (using LLM, this will take ~10s)") as status:
-                response = RestAPI(session.token).create(name=name, query=description)
+                response = RestAPI(get_current_token()).create(name=name, query=description)
                 code = convert_to_cli(response.code, session.config.repository.language, name)
                 prompt_path.parent.mkdir(parents=True, exist_ok=True)
                 prompt_path.write_text(response.context)

src/codegen/cli/commands/create/main.py

@@ -6,21 +6,21 @@
 
 from codegen.cli.api.client import RestAPI
 from codegen.cli.auth.decorators import requires_auth
-from codegen.cli.auth.session import CodegenSession
+from codegen.cli.auth.token_manager import get_current_token
 from codegen.cli.rich.codeblocks import format_command
 from codegen.cli.rich.spinners import create_spinner
 from codegen.cli.utils.codemod_manager import CodemodManager
 from codegen.cli.utils.function_finder import DecoratedFunction
 
 
-def deploy_functions(session: CodegenSession, functions: list[DecoratedFunction], message: str | None = None) -> None:
+def deploy_functions(functions: list[DecoratedFunction], message: str | None = None) -> None:
     """Deploy a list of functions."""
     if not functions:
         rich.print("\n[yellow]No @codegen.function decorators found.[/yellow]\n")
         return
 
     # Deploy each function
-    api_client = RestAPI(session.token)
+    api_client = RestAPI(get_current_token())
     rich.print()  # Add a blank line before deployments
 
     for func in functions:
@@ -47,7 +47,7 @@ def deploy_functions(session: CodegenSession, functions: list[DecoratedFunction]
 @click.argument("name", required=False)
 @click.option("-d", "--directory", type=click.Path(exists=True, path_type=Path), help="Directory to search for functions")
 @click.option("-m", "--message", help="Optional message to include with the deploy")
-def deploy_command(session: CodegenSession, name: str | None = None, directory: Path | None = None, message: str | None = None):
+def deploy_command(name: str | None = None, directory: Path | None = None, message: str | None = None):
     """Deploy codegen functions.
 
     If NAME is provided, deploys a specific function by that name.
@@ -70,11 +70,11 @@ def deploy_command(session: CodegenSession, name: str | None = None, directory:
                     rich.print(f"  • {func.filepath}")
                 msg = "Please specify the exact directory with --directory"
                 raise click.ClickException(msg)
-            deploy_functions(session, matching, message=message)
+            deploy_functions(matching, message=message)
         else:
             # Deploy all functions in the directory
             functions = CodemodManager.get_decorated(search_path)
-            deploy_functions(session, functions)
+            deploy_functions(functions)
     except Exception as e:
         msg = f"Failed to deploy: {e!s}"
         raise click.ClickException(msg)

src/codegen/cli/commands/deploy/main.py

@@ -4,7 +4,7 @@
 
 from codegen.cli.api.client import RestAPI
 from codegen.cli.auth.decorators import requires_auth
-from codegen.cli.auth.session import CodegenSession
+from codegen.cli.auth.token_manager import get_current_token
 from codegen.cli.errors import ServerError
 from codegen.cli.workspace.decorators import requires_init
 
@@ -13,13 +13,13 @@
 @click.option("--query", "-q", help="The question to ask the expert.")
 @requires_auth
 @requires_init
-def expert_command(session: CodegenSession, query: str):
+def expert_command(query: str):
     """Asks a codegen expert a question."""
     status = Status("Asking expert...", spinner="dots", spinner_style="purple")
     status.start()
 
     try:
-        response = RestAPI(session.token).ask_expert(query)
+        response = RestAPI(get_current_token()).ask_expert(query)
         status.stop()
         rich.print("[bold green]✓ Response received[/bold green]")
         rich.print(response.response)

src/codegen/cli/commands/expert/main.py

@@ -24,6 +24,7 @@ def init_command(path: str | None = None, token: str | None = None, language: st
     # Print a message if not in a git repo
     path = Path.cwd() if path is None else Path(path)
     repo_path = get_git_root_path(path)
+    rich.print(f"Found git repository at: {repo_path}")
     if repo_path is None:
         rich.print(f"\n[bold red]Error:[/bold red] Path={path} is not in a git repository")
         rich.print("[white]Please run this command from within a git repository.[/white]")

src/codegen/cli/commands/init/main.py

@@ -1,11 +1,7 @@
-import sys
-
-import rich
 import rich_click as click
 
-from codegen.cli.auth.auth_session import CodegenAuthenticatedSession
 from codegen.cli.auth.login import login_routine
-from codegen.cli.auth.token_manager import TokenManager, get_current_token
+from codegen.cli.auth.token_manager import get_current_token
 
 
 @click.command(name="login")
@@ -17,19 +13,4 @@ def login_command(token: str):
         msg = "Already authenticated. Use 'codegen logout' to clear the token."
         raise click.ClickException(msg)
 
-    if not token:
-        login_routine()
-        sys.exit(1)
-
-    # Use provided token or go through login flow
-    token_manager = TokenManager()
-    session = CodegenAuthenticatedSession(token=token)
-    try:
-        session.assert_authenticated()
-        token_manager.save_token(token)
-        rich.print(f"[green]✓ Stored token to:[/green] {token_manager.token_file}")
-        rich.print("[cyan]📊 Hey![/cyan] We collect anonymous usage data to improve your experience 🔒")
-        rich.print("To opt out, set [green]telemetry_enabled = false[/green] in [cyan]~/.config/codegen-sh/analytics.json[/cyan] ✨")
-    except ValueError as e:
-        msg = f"Error: {e!s}"
-        raise click.ClickException(msg)
+    login_routine(token)