fix: minor cleanup (#1217)

# Motivation

<!-- Why is this change necessary? -->

# Content

<!-- Please include a summary of the change -->

# Testing

<!-- How was the change tested? -->

# Please check the following before marking your PR as ready for review

- [ ] I have added tests for my changes
- [ ] I have updated the documentation or added new documentation as
needed

Author: rushilpatel0

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,44 +84,41 @@ 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
     """
     try:
         # 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)
+    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)

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,19 +53,15 @@ 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)}))
 
 
 if __name__ == "__main__":
     main()
-

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,32 +61,21 @@ 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))
 
     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()
+    main()

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()
-