Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
174 changes: 172 additions & 2 deletions src/codegen/cli/auth/token_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,15 @@ def save_token_with_org_info(self, token: str) -> None:
# Add organization info if available
orgs = org_data.get("items", [])
if orgs and len(orgs) > 0:
primary_org = orgs[0] # Use first org as primary
auth_data["organization"] = {"id": primary_org.get("id"), "name": primary_org.get("name"), "all_orgs": [{"id": org.get("id"), "name": org.get("name")} for org in orgs]}
# Store ALL organizations in cache for local resolution
all_orgs = [{"id": org.get("id"), "name": org.get("name")} for org in orgs]
primary_org = orgs[0] # Use first org as primary/default
auth_data["organization"] = {
"id": primary_org.get("id"),
"name": primary_org.get("name"),
"all_orgs": all_orgs
}
auth_data["organizations_cache"] = all_orgs # Separate cache for easy access

except requests.RequestException as e:
# If we can't fetch org info, still save the token but without org data
Expand Down Expand Up @@ -171,6 +178,53 @@ def get_user_info(self) -> dict | None:
return auth_data["user"]
return None

def get_cached_organizations(self) -> list[dict] | None:
"""Get all cached organizations.

Returns:
List of organization dictionaries with 'id' and 'name' keys, or None if no cache.
"""
auth_data = self.get_auth_data()
if auth_data and "organizations_cache" in auth_data:
return auth_data["organizations_cache"]
# Fallback to legacy format
if auth_data and "organization" in auth_data and "all_orgs" in auth_data["organization"]:
return auth_data["organization"]["all_orgs"]
return None

def is_org_id_in_cache(self, org_id: int) -> bool:
"""Check if an organization ID exists in the local cache.

Args:
org_id: The organization ID to check

Returns:
True if the organization ID is found in cache, False otherwise.
"""
cached_orgs = self.get_cached_organizations()
if not cached_orgs:
return False

return any(org.get("id") == org_id for org in cached_orgs)

def get_org_name_from_cache(self, org_id: int) -> str | None:
"""Get organization name from cache by ID.

Args:
org_id: The organization ID to look up

Returns:
Organization name if found in cache, None otherwise.
"""
cached_orgs = self.get_cached_organizations()
if not cached_orgs:
return None

for org in cached_orgs:
if org.get("id") == org_id:
return org.get("name")
return None


def get_current_token() -> str | None:
"""Get the current authentication token if one exists.
Expand Down Expand Up @@ -233,6 +287,42 @@ def get_current_org_name() -> str | None:
return token_manager.get_org_name()


def get_cached_organizations() -> list[dict] | None:
"""Get all cached organizations.

Returns:
List of organization dictionaries with 'id' and 'name' keys, or None if no cache.
"""
token_manager = TokenManager()
return token_manager.get_cached_organizations()


def is_org_id_cached(org_id: int) -> bool:
"""Check if an organization ID exists in the local cache.

Args:
org_id: The organization ID to check

Returns:
True if the organization ID is found in cache, False otherwise.
"""
token_manager = TokenManager()
return token_manager.is_org_id_in_cache(org_id)


def get_org_name_from_cache(org_id: int) -> str | None:
"""Get organization name from cache by ID.

Args:
org_id: The organization ID to look up

Returns:
Organization name if found in cache, None otherwise.
"""
token_manager = TokenManager()
return token_manager.get_org_name_from_cache(org_id)


def get_current_user_info() -> dict | None:
"""Get the stored user info if available.

Expand All @@ -241,3 +331,83 @@ def get_current_user_info() -> dict | None:
"""
token_manager = TokenManager()
return token_manager.get_user_info()


# Repository caching functions (similar to organization caching)

def get_cached_repositories() -> list[dict] | None:
"""Get all cached repositories.

Returns:
List of repository dictionaries with 'id' and 'name' keys, or None if no cache.
"""
token_manager = TokenManager()
auth_data = token_manager.get_auth_data()
if auth_data and "repositories_cache" in auth_data:
return auth_data["repositories_cache"]
return None


def cache_repositories(repositories: list[dict]) -> None:
"""Cache repositories to local storage.

Args:
repositories: List of repository dictionaries to cache
"""
token_manager = TokenManager()
auth_data = token_manager.get_auth_data()
if auth_data:
auth_data["repositories_cache"] = repositories
# Save back to file
try:
import json
with open(token_manager.token_file, 'w') as f:
json.dump(auth_data, f, indent=2)
except Exception:
pass # Fail silently


def is_repo_id_cached(repo_id: int) -> bool:
"""Check if a repository ID exists in the local cache.

Args:
repo_id: The repository ID to check

Returns:
True if the repository ID is found in cache, False otherwise.
"""
cached_repos = get_cached_repositories()
if not cached_repos:
return False

return any(repo.get("id") == repo_id for repo in cached_repos)


def get_repo_name_from_cache(repo_id: int) -> str | None:
"""Get repository name from cache by ID.

Args:
repo_id: The repository ID to look up

Returns:
Repository name if found in cache, None otherwise.
"""
cached_repos = get_cached_repositories()
if not cached_repos:
return None

for repo in cached_repos:
if repo.get("id") == repo_id:
return repo.get("name")

return None


def get_current_repo_name() -> str | None:
"""Get the current repository name from environment or cache."""
from codegen.cli.utils.repo import get_current_repo_id

repo_id = get_current_repo_id()
if repo_id:
return get_repo_name_from_cache(repo_id)
return None
4 changes: 4 additions & 0 deletions src/codegen/cli/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@
from codegen.cli.commands.integrations.main import integrations_app
from codegen.cli.commands.login.main import login
from codegen.cli.commands.logout.main import logout
from codegen.cli.commands.org.main import org
from codegen.cli.commands.profile.main import profile
from codegen.cli.commands.repo.main import repo
from codegen.cli.commands.style_debug.main import style_debug
from codegen.cli.commands.tools.main import tools
from codegen.cli.commands.tui.main import tui
Expand All @@ -39,7 +41,9 @@ def version_callback(value: bool):
main.command("init", help="Initialize or update the Codegen folder.")(init)
main.command("login", help="Store authentication token.")(login)
main.command("logout", help="Clear stored authentication token.")(logout)
main.command("org", help="Manage and switch between organizations.")(org)
main.command("profile", help="Display information about the currently authenticated user.")(profile)
main.command("repo", help="Manage repository configuration and environment variables.")(repo)
main.command("style-debug", help="Debug command to visualize CLI styling (spinners, etc).")(style_debug)
main.command("tools", help="List available tools from the Codegen API.")(tools)
main.command("tui", help="Launch the interactive TUI interface.")(tui)
Expand Down
119 changes: 92 additions & 27 deletions src/codegen/cli/commands/claude/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,36 +7,82 @@
import threading
import time

import requests
import typer
from rich import box
from rich.panel import Panel


from codegen.cli.api.endpoints import API_ENDPOINT
from codegen.cli.auth.token_manager import get_current_token
from codegen.cli.commands.claude.claude_log_watcher import ClaudeLogWatcherManager
from codegen.cli.commands.claude.claude_session_api import end_claude_session, generate_session_id
from codegen.cli.commands.claude.config.mcp_setup import add_codegen_mcp_server, cleanup_codegen_mcp_server
from codegen.cli.commands.claude.hooks import cleanup_claude_hook, ensure_claude_hook, get_codegen_url
from codegen.cli.commands.claude.quiet_console import console
from rich.console import Console

t_console = Console()

from codegen.cli.rich.spinners import create_spinner
from codegen.cli.utils.org import resolve_org_id


def claude(
org_id: int | None = typer.Option(None, help="Organization ID (defaults to CODEGEN_ORG_ID/REPOSITORY_ORG_ID or auto-detect)"),
no_mcp: bool | None = typer.Option(False, "--no-mcp", help="Disable Codegen's MCP server with additional capabilities over HTTP"),
):
"""Run Claude Code with session tracking.
def _run_claude_background(resolved_org_id: int, prompt: str | None) -> None:
"""Create a background agent run with Claude context and exit."""
token = get_current_token()
if not token:
console.print("[red]Error:[/red] Not authenticated. Please run 'codegen login' first.")
raise typer.Exit(1)

payload = {"prompt": prompt or "Start a Claude Code background session"}

This command runs Claude Code and tracks the session in the backend API:
- Generates a unique session ID
- Creates an agent run when Claude starts
- Updates the agent run status when Claude exits
"""
spinner = create_spinner("Creating agent run...")
spinner.start()
try:
headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"x-codegen-client": "codegen__claude_code",
}
url = f"{API_ENDPOINT.rstrip('/')}/v1/organizations/{resolved_org_id}/agent/run"
response = requests.post(url, headers=headers, json=payload)
response.raise_for_status()
agent_run_data = response.json()
finally:
spinner.stop()

run_id = agent_run_data.get("id", "Unknown")
status = agent_run_data.get("status", "Unknown")
web_url = agent_run_data.get("web_url", "")

result_lines = [
f"[cyan]Agent Run ID:[/cyan] {run_id}",
f"[cyan]Status:[/cyan] {status}",
]
if web_url:
result_lines.append(f"[cyan]Web URL:[/cyan] {web_url}")

t_console.print(
Panel(
"\n".join(result_lines),
title="🤖 [bold]Background Agent Run Created[/bold]",
border_style="green",
box=box.ROUNDED,
padding=(1, 2),
)
)
t_console.print("\n[dim]💡 Track progress with:[/dim] [cyan]codegen agents[/cyan]")
if web_url:
t_console.print(f"[dim]🌐 View in browser:[/dim] [link]{web_url}[/link]")


def _run_claude_interactive(resolved_org_id: int, no_mcp: bool | None) -> None:
"""Launch Claude Code with session tracking and log watching."""
# Generate session ID for tracking
session_id = generate_session_id()
console.print(f"🆔 Generated session ID: {session_id[:8]}...", style="dim")

# Resolve org_id early for session management
resolved_org_id = resolve_org_id(org_id)
if resolved_org_id is None:
console.print("[red]Error:[/red] Organization ID not provided. Pass --org-id, set CODEGEN_ORG_ID, or REPOSITORY_ORG_ID.")
raise typer.Exit(1)


console.print("🚀 Starting Claude Code with session tracking...", style="blue")
console.print(f"🎯 Organization ID: {resolved_org_id}", style="dim")

Expand Down Expand Up @@ -79,29 +125,27 @@ def claude(
url = get_codegen_url(session_id)
console.print(f"\n🔵 Codegen URL: {url}\n", style="bold blue")


process = subprocess.Popen(["claude", "--session-id", session_id])


# Start log watcher for the session
console.print("📋 Starting log watcher...", style="blue")
log_watcher_started = log_watcher_manager.start_watcher(
session_id=session_id,
org_id=resolved_org_id,
poll_interval=1.0, # Check every second
on_log_entry=None
poll_interval=1.0,
on_log_entry=None,
)

if not log_watcher_started:
console.print("⚠️ Failed to start log watcher", style="yellow")

# Handle Ctrl+C gracefully
def signal_handler(signum, frame):
console.print("\n🛑 Stopping Claude Code...", style="yellow")
log_watcher_manager.stop_all_watchers() # Stop log watchers
log_watcher_manager.stop_all_watchers()
process.terminate()
cleanup_claude_hook() # Clean up our hook
cleanup_codegen_mcp_server() # Clean up MCP Server
cleanup_claude_hook()
cleanup_codegen_mcp_server()
end_claude_session(session_id, "ERROR", resolved_org_id)
sys.exit(0)

Expand Down Expand Up @@ -140,12 +184,33 @@ def signal_handler(signum, frame):
log_watcher_manager.stop_all_watchers()
except Exception as e:
console.print(f"⚠️ Error stopping log watchers: {e}", style="yellow")

cleanup_claude_hook()

# Show final session info
url = get_codegen_url(session_id)
console.print(f"\n🔵 Session URL: {url}", style="bold blue")
console.print(f"🆔 Session ID: {session_id}", style="dim")
console.print(f"🎯 Organization ID: {resolved_org_id}", style="dim")
console.print("💡 Check your backend to see the session data", style="dim")
console.print("💡 Check your backend to see the session data", style="dim")


def claude(
org_id: int | None = typer.Option(None, help="Organization ID (defaults to CODEGEN_ORG_ID/REPOSITORY_ORG_ID or auto-detect)"),
no_mcp: bool | None = typer.Option(False, "--no-mcp", help="Disable Codegen's MCP server with additional capabilities over HTTP"),
background: str | None = typer.Option(None, "--background", "-b", help="Create a background agent run with this prompt instead of launching Claude Code"),
):
"""Run Claude Code with session tracking or create a background run."""
# Resolve org_id early for session management
resolved_org_id = resolve_org_id(org_id)
if resolved_org_id is None:
console.print("[red]Error:[/red] Organization ID not provided. Pass --org-id, set CODEGEN_ORG_ID, or REPOSITORY_ORG_ID.")
raise typer.Exit(1)

if background is not None:
# Use the value from --background as the prompt, with --prompt as fallback
final_prompt = background or prompt
_run_claude_background(resolved_org_id, final_prompt)
return

_run_claude_interactive(resolved_org_id, no_mcp)
5 changes: 5 additions & 0 deletions src/codegen/cli/commands/org/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""Organization management command."""

from .main import org

__all__ = ["org"]
Loading
Loading