feat: Validate session on create and fetch (#471)

Author: caroljung-cg

src/codegen/cli/auth/decorators.py

@@ -19,7 +19,7 @@ def wrapper(*args, **kwargs):
         session = CodegenSession.from_active_session()
 
         # Check for valid session
-        if session is None or not session.is_valid():
+        if session is None:
             pretty_print_error("There is currently no active session.\nPlease run 'codegen init' to initialize the project.")
             raise click.Abort()
 

src/codegen/cli/auth/session.py

@@ -1,8 +1,13 @@
 from pathlib import Path
 
-from pygit2.repository import Repository
+import click
+import rich
+from github import BadCredentialsException
+from github.MainClass import Github
 
 from codegen.cli.git.repo import get_git_repo
+from codegen.cli.rich.codeblocks import format_command
+from codegen.git.repo_operator.local_git_repo import LocalGitRepo
 from codegen.shared.configs.constants import CODEGEN_DIR_NAME, CONFIG_FILENAME
 from codegen.shared.configs.models.session import SessionConfig
 from codegen.shared.configs.session_configs import global_config, load_session_config
@@ -11,16 +16,25 @@
 class CodegenSession:
     """Represents an authenticated codegen session with user and repository context"""
 
-    repo_path: Path  # TODO: rename to root_path
+    repo_path: Path
+    local_git: LocalGitRepo
     codegen_dir: Path
     config: SessionConfig
     existing: bool
 
-    def __init__(self, repo_path: Path):
+    def __init__(self, repo_path: Path, git_token: str | None = None) -> None:
+        if not repo_path.exists() or get_git_repo(repo_path) is None:
+            rich.print(f"\n[bold red]Error:[/bold red] Path to git repo does not exist at {self.repo_path}")
+            raise click.Abort()
+
         self.repo_path = repo_path
+        self.local_git = LocalGitRepo(repo_path=repo_path)
         self.codegen_dir = repo_path / CODEGEN_DIR_NAME
-        self.existing = global_config.get_session(repo_path) is not None
         self.config = load_session_config(self.codegen_dir / CONFIG_FILENAME)
+        self.config.secrets.github_token = git_token or self.config.secrets.github_token
+        self.existing = global_config.get_session(repo_path) is not None
+
+        self._initialize()
         global_config.set_active_session(repo_path)
 
     @classmethod
@@ -31,19 +45,52 @@ def from_active_session(cls) -> "CodegenSession | None":
 
         return cls(active_session)
 
-    def is_valid(self) -> bool:
-        """Validates that the session configuration is correct"""
-        # TODO: also make sure all the expected prompt, jupyter, codemods are present
-        # TODO: make sure there is still a git instance here.
-        return self.repo_path.exists() and self.codegen_dir.exists() and Path(self.config.file_path).exists()
-
-    @property
-    def git_repo(self) -> Repository:
-        git_repo = get_git_repo(Path.cwd())
-        if not git_repo:
-            msg = "No git repository found"
-            raise ValueError(msg)
-        return git_repo
+    def _initialize(self) -> None:
+        """Initialize the codegen session"""
+        self._validate()
+
+        self.config.repository.repo_path = self.config.repository.repo_path or str(self.local_git.repo_path)
+        self.config.repository.repo_name = self.config.repository.repo_name or self.local_git.name
+        self.config.repository.full_name = self.config.repository.full_name or self.local_git.full_name
+        self.config.repository.user_name = self.config.repository.user_name or self.local_git.user_name
+        self.config.repository.user_email = self.config.repository.user_email or self.local_git.user_email
+        self.config.repository.language = self.config.repository.language or self.local_git.get_language(access_token=self.config.secrets.github_token).upper()
+        self.config.save()
+
+    def _validate(self) -> None:
+        """Validates that the session configuration is correct, otherwise raises an error"""
+        if not self.codegen_dir.exists():
+            rich.print(f"\n[bold red]Error:[/bold red] Codegen folder is missing at {self.codegen_dir}")
+            raise click.Abort()
+
+        if not Path(self.config.file_path).exists():
+            rich.print(f"\n[bold red]Error:[/bold red] Missing config.toml at {self.codegen_dir}")
+            rich.print("[white]Please remove the codegen folder and reinitialize.[/white]")
+            rich.print(format_command(f"rm -rf {self.codegen_dir} && codegen init"))
+            raise click.Abort()
+
+        git_token = self.config.secrets.github_token
+        if git_token is None:
+            rich.print("\n[bold yellow]Warning:[/bold yellow] GitHub token not found")
+            rich.print("To enable full functionality, please set your GitHub token:")
+            rich.print(format_command("export CODEGEN_SECRETS__GITHUB_TOKEN=<your-token>"))
+            rich.print("Or pass in as a parameter:")
+            rich.print(format_command("codegen init --token <your-token>"))
+            raise click.Abort()
+
+        if self.local_git.origin_remote is None:
+            rich.print("\n[bold red]Error:[/bold red] No remote found for repository")
+            rich.print("[white]Please add a remote to the repository.[/white]")
+            rich.print("\n[dim]To add a remote to the repository:[/dim]")
+            rich.print(format_command("git remote add origin <your-repo-url>"))
+            raise click.Abort()
+
+        try:
+            Github(login_or_token=git_token).get_repo(self.local_git.full_name)
+        except BadCredentialsException:
+            rich.print(format_command(f"\n[bold red]Error:[/bold red] Invalid GitHub token={git_token} for repo={self.local_git.full_name}"))
+            rich.print("[white]Please provide a valid GitHub token for this repository.[/white]")
+            raise click.Abort()
 
     def __str__(self) -> str:
         return f"CodegenSession(user={self.config.repository.user_name}, repo={self.config.repository.repo_name})"

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

@@ -3,14 +3,11 @@
 
 import rich
 import rich_click as click
-from github import BadCredentialsException
-from github.MainClass import Github
 
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.commands.init.render import get_success_message
 from codegen.cli.rich.codeblocks import format_command
 from codegen.cli.workspace.initialize_workspace import initialize_codegen
-from codegen.git.repo_operator.local_git_repo import LocalGitRepo
 from codegen.git.utils.path import get_git_root_path
 
 
@@ -25,6 +22,7 @@ def init_command(path: str | None = None, token: str | None = None, language: st
     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]")
@@ -33,45 +31,10 @@ def init_command(path: str | None = None, token: str | None = None, language: st
         rich.print(format_command("codegen init"))
         sys.exit(1)
 
-    local_git = LocalGitRepo(repo_path=repo_path)
-    if local_git.origin_remote is None:
-        rich.print("\n[bold red]Error:[/bold red] No remote found for repository")
-        rich.print("[white]Please add a remote to the repository.[/white]")
-        rich.print("\n[dim]To add a remote to the repository:[/dim]")
-        rich.print(format_command("git remote add origin <your-repo-url>"))
-        sys.exit(1)
-
-    session = CodegenSession(repo_path=repo_path)
-    config = session.config
-
-    if token is None:
-        token = config.secrets.github_token
-    else:
-        config.secrets.github_token = token
-
-    if not token:
-        rich.print("\n[bold yellow]Warning:[/bold yellow] GitHub token not found")
-        rich.print("To enable full functionality, please set your GitHub token:")
-        rich.print(format_command("export CODEGEN_SECRETS__GITHUB_TOKEN=<your-token>"))
-        rich.print("Or pass in as a parameter:")
-        rich.print(format_command("codegen init --token <your-token>"))
-    else:
-        # Validate that the token is valid for the repo path.
-        try:
-            Github(login_or_token=token).get_repo(local_git.full_name)
-        except BadCredentialsException:
-            rich.print(format_command(f"\n[bold red]Error:[/bold red] Invalid GitHub token={token} for repo={local_git.full_name}"))
-            rich.print("[white]Please provide a valid GitHub token for this repository.[/white]")
-            sys.exit(1)
-
-    # Save repo config
-    config.repository.repo_path = str(local_git.repo_path)
-    config.repository.repo_name = local_git.name
-    config.repository.full_name = local_git.full_name
-    config.repository.user_name = local_git.user_name
-    config.repository.user_email = local_git.user_email
-    config.repository.language = (language or local_git.get_language(access_token=token)).upper()
-    config.save()
+    session = CodegenSession(repo_path=repo_path, git_token=token)
+    if language:
+        session.config.repository.language = language.upper()
+        session.config.save()
 
     action = "Updating" if session.existing else "Initializing"
     codegen_dir, docs_dir, examples_dir = initialize_codegen(status=action, session=session, fetch_docs=fetch_docs)

src/codegen/cli/commands/run/run_cloud.py

@@ -2,6 +2,7 @@
 
 import rich
 import rich_click as click
+from pygit2._pygit2 import Repository
 from rich.panel import Panel
 
 from codegen.cli.api.client import RestAPI
@@ -87,7 +88,8 @@ def run_cloud(session: CodegenSession, function, apply_local: bool = False, diff
 
             if apply_local and run_output.observation:
                 try:
-                    apply_patch(session.git_repo, f"\n{run_output.observation}\n")
+                    git_repo = Repository(str(session.repo_path))
+                    apply_patch(git_repo, f"\n{run_output.observation}\n")
                     rich.print("")
                     rich.print("[green]✓ Changes have been applied to your local filesystem[/green]")
                     rich.print("[yellow]→ Don't forget to commit your changes:[/yellow]")

src/codegen/cli/git/url.py

@@ -2,8 +2,6 @@
 import sys
 from collections.abc import Callable
 
-import click
-
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.rich.pretty_print import pretty_print_error
 
@@ -19,11 +17,6 @@ def wrapper(*args, **kwargs):
             pretty_print_error("Codegen not initialized. Please run `codegen init` from a git repo workspace.")
             sys.exit(1)
 
-        # 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.")
-            raise click.Abort()
-
         kwargs["session"] = session
         return f(*args, **kwargs)
 

src/codegen/cli/workspace/decorators.py

@@ -4,15 +4,13 @@
 
 import requests
 import rich
-import toml
 from rich.status import Status
 
 from codegen.cli.api.client import RestAPI
+from codegen.cli.api.endpoints import CODEGEN_SYSTEM_PROMPT_URL
 from codegen.cli.auth.constants import CODEGEN_DIR, DOCS_DIR, EXAMPLES_DIR, PROMPTS_DIR
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.auth.token_manager import get_current_token
-from codegen.cli.git.repo import get_git_repo
-from codegen.cli.git.url import get_git_organization_and_repo
 from codegen.cli.rich.spinners import create_spinner
 from codegen.cli.utils.notebooks import create_notebook
 from codegen.cli.workspace.docs_workspace import populate_api_docs
@@ -31,12 +29,10 @@ def initialize_codegen(session: CodegenSession, status: Status | str = "Initiali
     Returns:
         Tuple of (codegen_folder, docs_folder, examples_folder)
     """
-    repo = get_git_repo()
     CODEGEN_FOLDER = session.repo_path / CODEGEN_DIR
     PROMPTS_FOLDER = session.repo_path / PROMPTS_DIR
     DOCS_FOLDER = session.repo_path / DOCS_DIR
     EXAMPLES_FOLDER = session.repo_path / EXAMPLES_DIR
-    CONFIG_PATH = CODEGEN_FOLDER / "config.toml"
     JUPYTER_DIR = CODEGEN_FOLDER / "jupyter"
     CODEMODS_DIR = CODEGEN_FOLDER / "codemods"
     SYSTEM_PROMPT_PATH = CODEGEN_FOLDER / "codegen-system-prompt.txt"
@@ -62,35 +58,17 @@ def initialize_codegen(session: CodegenSession, status: Status | str = "Initiali
 
         # Download system prompt
         try:
-            from codegen.cli.api.endpoints import CODEGEN_SYSTEM_PROMPT_URL
-
             response = requests.get(CODEGEN_SYSTEM_PROMPT_URL)
             response.raise_for_status()
             SYSTEM_PROMPT_PATH.write_text(response.text)
         except Exception as e:
             rich.print(f"[yellow]Warning: Could not download system prompt: {e}[/yellow]")
 
-        if not repo:
-            rich.print("No git repository found. Please run this command in a git repository.")
-        else:
-            status_obj.update(f"   {'Updating' if isinstance(status, Status) else status} .gitignore...")
-            modify_gitignore(CODEGEN_FOLDER)
-
-            # Create or update config.toml with basic repo info
-            org_name, repo_name = get_git_organization_and_repo(repo)
-            config = {}
-            if CONFIG_PATH.exists():
-                config = toml.load(CONFIG_PATH)
-            config.update(
-                {
-                    "organization_name": config.get("organization_name", org_name),
-                    "repo_name": config.get("repo_name", repo_name),
-                }
-            )
-            CONFIG_PATH.write_text(toml.dumps(config))
-
-            # Create notebook template
-            create_notebook(JUPYTER_DIR)
+        status_obj.update(f"   {'Updating' if isinstance(status, Status) else status} .gitignore...")
+        modify_gitignore(CODEGEN_FOLDER)
+
+        # Create notebook template
+        create_notebook(JUPYTER_DIR)
 
         # Only fetch docs and examples if requested and session is provided
         if fetch_docs and session:

src/codegen/cli/workspace/initialize_workspace.py

@@ -2,6 +2,7 @@
 from functools import cached_property
 from pathlib import Path
 
+import giturlparse
 from git import Repo
 from git.remote import Remote
 
@@ -30,8 +31,8 @@ def full_name(self) -> str | None:
         if not self.origin_remote:
             return None
 
-        url_segments = self.origin_remote.url.split("/")
-        return f"{url_segments[-2]}/{url_segments[-1].replace('.git', '')}"
+        parsed = giturlparse.parse(self.origin_remote.url)
+        return f"{parsed.owner}/{parsed.name}"
 
     @cached_property
     def origin_remote(self) -> Remote | None: