From ebdef158acc818b84874a25201821a87fe528c3e Mon Sep 17 00:00:00 2001 From: Rushil Patel Date: Tue, 19 Aug 2025 11:18:11 -0700 Subject: [PATCH] fix: minor cleanup --- .../cli/commands/claude/claude_session_api.py | 107 ++++++++---------- .../config/claude_session_active_hook.py | 14 +-- .../claude/config/claude_session_hook.py | 33 ++---- .../claude/config/claude_session_stop_hook.py | 14 +-- src/codegen/cli/commands/claude/hooks.py | 8 +- src/codegen/cli/commands/claude/main.py | 12 +- 6 files changed, 75 insertions(+), 113 deletions(-) diff --git a/src/codegen/cli/commands/claude/claude_session_api.py b/src/codegen/cli/commands/claude/claude_session_api.py index a931f0561..6ff5e5566 100644 --- a/src/codegen/cli/commands/claude/claude_session_api.py +++ b/src/codegen/cli/commands/claude/claude_session_api.py @@ -12,9 +12,9 @@ from codegen.cli.utils.org import resolve_org_id - class ClaudeSessionAPIError(Exception): """Exception raised for Claude session API errors.""" + pass @@ -25,14 +25,14 @@ def generate_session_id() -> str: def create_claude_session(session_id: str, org_id: Optional[int] = None) -> Optional[str]: """Create a new Claude Code session in the backend. - + Args: session_id: The session ID to register org_id: Organization ID (will be resolved if None) - + Returns: Agent run ID if successful, None if failed - + Raises: ClaudeSessionAPIError: If the API call fails """ @@ -42,24 +42,21 @@ def create_claude_session(session_id: str, org_id: Optional[int] = None) -> Opti if resolved_org_id is None: console.print("⚠️ Could not resolve organization ID for session creation", style="yellow") return None - + # Get authentication token token = get_current_token() if not token: console.print("⚠️ No authentication token found for session creation", style="yellow") return None - + # Prepare API request url = f"{API_ENDPOINT.rstrip('/')}/v1/organizations/{resolved_org_id}/claude_code/session" - headers = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json" - } + headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} payload = {"session_id": session_id} - + # Make API request response = requests.post(url, json=payload, headers=headers, timeout=30) - + if response.status_code == 200: try: result = response.json() @@ -75,10 +72,10 @@ def create_claude_session(session_id: str, org_id: Optional[int] = None) -> Opti error_msg = f"{error_msg}: {error_detail}" except Exception: error_msg = f"{error_msg}: {response.text}" - + console.print(f"⚠️ Failed to create Claude session: {error_msg}", style="yellow") return None - + except requests.RequestException as e: console.print(f"⚠️ Network error creating Claude session: {e}", style="yellow") return None @@ -87,14 +84,14 @@ def create_claude_session(session_id: str, org_id: Optional[int] = None) -> Opti return None -def end_claude_session(session_id: str, status: str, org_id: Optional[int] = None) -> bool: - """End a Claude Code session in the backend. - +def update_claude_session_status(session_id: str, status: str, org_id: Optional[int] = None) -> bool: + """Update a Claude Code session status in the backend. + Args: - session_id: The session ID to end - status: Completion status ("COMPLETE" or "ERROR") + session_id: The session ID to update + status: Session status ("COMPLETE", "ERROR", "ACTIVE", etc.) org_id: Organization ID (will be resolved if None) - + Returns: True if successful, False if failed """ @@ -102,29 +99,26 @@ def end_claude_session(session_id: str, status: str, org_id: Optional[int] = Non # Resolve org_id resolved_org_id = resolve_org_id(org_id) if resolved_org_id is None: - console.print("⚠️ Could not resolve organization ID for session completion", style="yellow") + console.print("⚠️ Could not resolve organization ID for session status update", style="yellow") return False - + # Get authentication token token = get_current_token() if not token: - console.print("⚠️ No authentication token found for session completion", style="yellow") + console.print("⚠️ No authentication token found for session status update", style="yellow") return False - + # Prepare API request url = f"{API_ENDPOINT.rstrip('/')}/v1/organizations/{resolved_org_id}/claude_code/session/{session_id}/status" - headers = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json" - } + headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} payload = {"status": status} - + # Make API request response = requests.post(url, json=payload, headers=headers, timeout=30) - + if response.status_code == 200: - status_emoji = "✅" if status == "COMPLETE" else "❌" - console.print(f"{status_emoji} Ended Claude session {session_id[:8]}... with status {status}", style="green") + status_emoji = "✅" if status == "COMPLETE" else "🔄" if status == "ACTIVE" else "❌" + console.print(f"{status_emoji} Updated Claude session {session_id[:8]}... status to {status}", style="green") return True else: error_msg = f"HTTP {response.status_code}" @@ -133,26 +127,26 @@ def end_claude_session(session_id: str, status: str, org_id: Optional[int] = Non error_msg = f"{error_msg}: {error_detail}" except Exception: error_msg = f"{error_msg}: {response.text}" - - console.print(f"⚠️ Failed to end Claude session: {error_msg}", style="yellow") + + console.print(f"⚠️ Failed to update Claude session status: {error_msg}", style="yellow") return False - + except requests.RequestException as e: - console.print(f"⚠️ Network error ending Claude session: {e}", style="yellow") + console.print(f"⚠️ Network error updating Claude session status: {e}", style="yellow") return False except Exception as e: - console.print(f"⚠️ Unexpected error ending Claude session: {e}", style="yellow") + console.print(f"⚠️ Unexpected error updating Claude session status: {e}", style="yellow") return False def send_claude_session_log(session_id: str, log_entry: dict, org_id: Optional[int] = None) -> bool: """Send a log entry to the Claude Code session log endpoint. - + Args: session_id: The session ID log_entry: The log entry to send (dict) org_id: Organization ID (will be resolved if None) - + Returns: True if successful, False if failed """ @@ -162,24 +156,21 @@ def send_claude_session_log(session_id: str, log_entry: dict, org_id: Optional[i if resolved_org_id is None: console.print("⚠️ Could not resolve organization ID for log sending", style="yellow") return False - + # Get authentication token token = get_current_token() if not token: console.print("⚠️ No authentication token found for log sending", style="yellow") return False - + # Prepare API request url = f"{API_ENDPOINT.rstrip('/')}/v1/organizations/{resolved_org_id}/claude_code/session/{session_id}/log" - headers = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json" - } + headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} payload = {"log": log_entry} - + # Make API request response = requests.post(url, json=payload, headers=headers, timeout=30) - + if response.status_code == 200: return True else: @@ -189,10 +180,10 @@ def send_claude_session_log(session_id: str, log_entry: dict, org_id: Optional[i error_msg = f"{error_msg}: {error_detail}" except Exception: error_msg = f"{error_msg}: {response.text}" - + console.print(f"⚠️ Failed to send log entry: {error_msg}", style="yellow") return False - + except requests.RequestException as e: console.print(f"⚠️ Network error sending log entry: {e}", style="yellow") return False @@ -203,25 +194,21 @@ def send_claude_session_log(session_id: str, log_entry: dict, org_id: Optional[i def write_session_hook_data(session_id: str, org_id: Optional[int] = None) -> str: """Write session data for Claude hook and create session via API. - + This function is called by the Claude hook to both write session data locally and create the session in the backend API. - + Args: session_id: The session ID org_id: Organization ID - + Returns: JSON string to write to the session file """ # Create session in backend API agent_run_id = create_claude_session(session_id, org_id) - + # Prepare session data - session_data = { - "session_id": session_id, - "agent_run_id": agent_run_id, - "org_id": resolve_org_id(org_id) - } - - return json.dumps(session_data, indent=2) \ No newline at end of file + session_data = {"session_id": session_id, "agent_run_id": agent_run_id, "org_id": resolve_org_id(org_id)} + + return json.dumps(session_data, indent=2) diff --git a/src/codegen/cli/commands/claude/config/claude_session_active_hook.py b/src/codegen/cli/commands/claude/config/claude_session_active_hook.py index 23404ba68..7aea27662 100644 --- a/src/codegen/cli/commands/claude/config/claude_session_active_hook.py +++ b/src/codegen/cli/commands/claude/config/claude_session_active_hook.py @@ -17,9 +17,9 @@ sys.path.insert(0, str(codegen_cli_dir)) try: - from codegen.cli.commands.claude.claude_session_api import end_claude_session + from codegen.cli.commands.claude.claude_session_api import update_claude_session_status except ImportError: - end_claude_session = None + update_claude_session_status = None def read_session_file() -> dict: @@ -53,14 +53,11 @@ def main(): except ValueError: org_id = None - if end_claude_session and session_id: - end_claude_session(session_id, "ACTIVE", org_id) + if update_claude_session_status and session_id: + update_claude_session_status(session_id, "ACTIVE", org_id) # Print minimal output - print(json.dumps({ - "session_id": session_id, - "status": "ACTIVE" - })) + print(json.dumps({"session_id": session_id, "status": "ACTIVE"})) except Exception as e: print(json.dumps({"error": str(e)})) @@ -68,4 +65,3 @@ def main(): if __name__ == "__main__": main() - diff --git a/src/codegen/cli/commands/claude/config/claude_session_hook.py b/src/codegen/cli/commands/claude/config/claude_session_hook.py index d72dde9ab..374ef0462 100755 --- a/src/codegen/cli/commands/claude/config/claude_session_hook.py +++ b/src/codegen/cli/commands/claude/config/claude_session_hook.py @@ -11,19 +11,14 @@ import sys from pathlib import Path +from codegen.cli.commands.claude.claude_session_api import create_claude_session +from codegen.cli.utils.org import resolve_org_id + # Add the codegen CLI to the path so we can import from it script_dir = Path(__file__).parent codegen_cli_dir = script_dir.parent.parent.parent sys.path.insert(0, str(codegen_cli_dir)) -try: - from codegen.cli.commands.claude.claude_session_api import create_claude_session - from codegen.cli.utils.org import resolve_org_id -except ImportError: - # Fallback if imports fail - just write basic session data - create_claude_session = None - resolve_org_id = None - def main(): """Main hook function called by Claude Code.""" @@ -44,10 +39,11 @@ def main(): if not session_id: # Fallback: try to extract from input data session_id = input_data.get("session_id") - + if not session_id: # Generate a basic session ID if none available import uuid + session_id = str(uuid.uuid4()) # Get org_id from environment variable (set by main.py) @@ -65,17 +61,11 @@ def main(): # Create session via API if available agent_run_id = None - if create_claude_session and org_id: + if org_id: agent_run_id = create_claude_session(session_id, org_id) # Prepare session data - session_data = { - "session_id": session_id, - "agent_run_id": agent_run_id, - "org_id": org_id, - "hook_event": input_data.get("hook_event_name"), - "timestamp": input_data.get("timestamp") - } + session_data = {"session_id": session_id, "agent_run_id": agent_run_id, "org_id": org_id, "hook_event": input_data.get("hook_event_name"), "timestamp": input_data.get("timestamp")} # Output the session data (this gets written to the session file by the hook command) print(json.dumps(session_data, indent=2)) @@ -83,14 +73,9 @@ def main(): except Exception as e: # If anything fails, at least output basic session data session_id = os.environ.get("CODEGEN_CLAUDE_SESSION_ID", "unknown") - fallback_data = { - "session_id": session_id, - "error": str(e), - "agent_run_id": None, - "org_id": None - } + fallback_data = {"session_id": session_id, "error": str(e), "agent_run_id": None, "org_id": None} print(json.dumps(fallback_data, indent=2)) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/codegen/cli/commands/claude/config/claude_session_stop_hook.py b/src/codegen/cli/commands/claude/config/claude_session_stop_hook.py index 9ea1e03ef..6e954d55c 100644 --- a/src/codegen/cli/commands/claude/config/claude_session_stop_hook.py +++ b/src/codegen/cli/commands/claude/config/claude_session_stop_hook.py @@ -17,9 +17,9 @@ sys.path.insert(0, str(codegen_cli_dir)) try: - from codegen.cli.commands.claude.claude_session_api import end_claude_session + from codegen.cli.commands.claude.claude_session_api import update_claude_session_status except ImportError: - end_claude_session = None + update_claude_session_status = None def read_session_file() -> dict: @@ -53,14 +53,11 @@ def main(): except ValueError: org_id = None - if end_claude_session and session_id: - end_claude_session(session_id, "COMPLETE", org_id) + if update_claude_session_status and session_id: + update_claude_session_status(session_id, "COMPLETE", org_id) # Print minimal output to avoid noisy hooks - print(json.dumps({ - "session_id": session_id, - "status": "COMPLETE" - })) + print(json.dumps({"session_id": session_id, "status": "COMPLETE"})) except Exception as e: # Ensure hook doesn't fail Claude if something goes wrong @@ -69,4 +66,3 @@ def main(): if __name__ == "__main__": main() - diff --git a/src/codegen/cli/commands/claude/hooks.py b/src/codegen/cli/commands/claude/hooks.py index e0d82b969..ac993d20b 100644 --- a/src/codegen/cli/commands/claude/hooks.py +++ b/src/codegen/cli/commands/claude/hooks.py @@ -2,12 +2,10 @@ import json import os -import time from pathlib import Path from codegen.cli.commands.claude.quiet_console import console - CLAUDE_CONFIG_DIR = Path.home() / ".claude" HOOKS_CONFIG_FILE = CLAUDE_CONFIG_DIR / "settings.json" CODEGEN_DIR = Path.home() / ".codegen" @@ -39,7 +37,7 @@ def ensure_claude_hook() -> bool: # Build the shell command that will create session via API and write session data start_hook_script_path = Path(__file__).parent / "config" / "claude_session_hook.py" - start_hook_command = f"mkdir -p {CODEGEN_DIR} && python3 {start_hook_script_path} > {SESSION_FILE}" + start_hook_command = f"python3 {start_hook_script_path} > {SESSION_FILE}" # Build the stop hook command to mark session COMPLETE stop_hook_script_path = Path(__file__).parent / "config" / "claude_session_stop_hook.py" @@ -99,9 +97,9 @@ def ensure_claude_hook() -> bool: console.print("✅ Replaced existing Claude hooks (SessionStart, Stop)", style="green") else: console.print("✅ Registered new Claude hooks (SessionStart, Stop)", style="green") - console.print(f" Start hook: {start_hook_command[:50]}...", style="dim") + console.print(f" Start hook: {start_hook_command}", style="dim") console.print(f" Stop hook: {stop_hook_command}", style="dim") - console.print(f" Active hook:{' ' if len('Active hook:')<1 else ''} {active_hook_command}", style="dim") + console.print(f" Active hook:{' ' if len('Active hook:') < 1 else ''} {active_hook_command}", style="dim") # Verify the hook was written correctly try: diff --git a/src/codegen/cli/commands/claude/main.py b/src/codegen/cli/commands/claude/main.py index a8f3a33c5..693a6a712 100644 --- a/src/codegen/cli/commands/claude/main.py +++ b/src/codegen/cli/commands/claude/main.py @@ -14,7 +14,7 @@ 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.claude_session_api import update_claude_session_status, 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 @@ -142,7 +142,7 @@ def signal_handler(signum, frame): process.terminate() cleanup_claude_hook() cleanup_codegen_mcp_server() - end_claude_session(session_id, "ERROR", resolved_org_id) + update_claude_session_status(session_id, "ERROR", resolved_org_id) sys.exit(0) signal.signal(signal.SIGINT, signal_handler) @@ -152,7 +152,7 @@ def signal_handler(signum, frame): # Handle session completion based on exit code session_status = "COMPLETE" if returncode == 0 else "ERROR" - end_claude_session(session_id, session_status, resolved_org_id) + update_claude_session_status(session_id, session_status, resolved_org_id) if returncode != 0: console.print(f"❌ Claude Code exited with error code {returncode}", style="red") @@ -163,16 +163,16 @@ def signal_handler(signum, frame): console.print("❌ Claude Code not found. Please install Claude Code first.", style="red") console.print("💡 Visit: https://claude.ai/download", style="dim") log_watcher_manager.stop_all_watchers() - end_claude_session(session_id, "ERROR", resolved_org_id) + update_claude_session_status(session_id, "ERROR", resolved_org_id) raise typer.Exit(1) except KeyboardInterrupt: console.print("\n🛑 Interrupted by user", style="yellow") log_watcher_manager.stop_all_watchers() - end_claude_session(session_id, "ERROR", resolved_org_id) + update_claude_session_status(session_id, "ERROR", resolved_org_id) except Exception as e: console.print(f"❌ Error running Claude Code: {e}", style="red") log_watcher_manager.stop_all_watchers() - end_claude_session(session_id, "ERROR", resolved_org_id) + update_claude_session_status(session_id, "ERROR", resolved_org_id) raise typer.Exit(1) finally: # Clean up resources