prevent oos errors

Author: eyw520

servers/fai-lambda/fai-scribe/src/handler.py

@@ -1,9 +1,11 @@
 import asyncio
 import json
+import shutil
 from datetime import (
     UTC,
     datetime,
 )
+from pathlib import Path
 from typing import Any
 
 import httpx
@@ -119,6 +121,14 @@ async def handle_editing_request(
                     f"{cleanup_response.status_code} - {cleanup_response.text}"
                 )
 
+            try:
+                session_dir = Path(repo_path).parent
+                if session_dir.exists() and session_dir.name.startswith("editing-"):
+                    shutil.rmtree(session_dir)
+                    LOGGER.info(f"Cleaned up session directory: {session_dir}")
+            except Exception as cleanup_error:
+                LOGGER.warning(f"Failed to cleanup session directory: {cleanup_error}")
+
             return {
                 "editing_id": editing_id,
                 "session_id": session_id,

servers/fai-lambda/shared/utils/git.py

@@ -2,13 +2,34 @@
 import os
 import shutil
 import subprocess
+import time
 from pathlib import Path
 
 logger = logging.getLogger()
 
 
+def cleanup_old_sessions(base_path: Path, max_age_hours: int = 24) -> None:
+    """Remove session directories older than max_age_hours."""
+    try:
+        if not base_path.exists():
+            return
+
+        cutoff_time = time.time() - (max_age_hours * 3600)
+
+        for item in base_path.iterdir():
+            if item.is_dir() and (item.name.startswith("editing-") or item.name.startswith("session-")):
+                try:
+                    if item.stat().st_mtime < cutoff_time:
+                        logger.info(f"Removing old session directory: {item.name}")
+                        shutil.rmtree(item)
+                except Exception as e:
+                    logger.warning(f"Failed to remove old session {item.name}: {e}")
+    except Exception as e:
+        logger.warning(f"Failed to cleanup old sessions: {e}")
+
+
 def clone_repo(repository: str, session_id: str, session_type: str = "session") -> str:
-    """Clone a GitHub repository into /tmp.
+    """Clone a GitHub repository into EFS storage.
 
     Args:
         repository: GitHub repository in format 'owner/repo'
@@ -22,36 +43,75 @@ def clone_repo(repository: str, session_id: str, session_type: str = "session")
     if not github_token:
         raise RuntimeError("GITHUB_TOKEN environment variable not set")
 
-    repo_path = Path("/tmp") / f"{session_type}-{session_id}" / repository
-    repo_path.parent.mkdir(parents=True, exist_ok=True)
+    base_path = Path(os.environ.get("HOME", "/mnt/efs")) / "repos"
+    cache_path = base_path / repository.replace("/", "_")
+    session_path = base_path / f"{session_type}-{session_id}" / repository
+
+    base_path.mkdir(parents=True, exist_ok=True)
+    session_path.parent.mkdir(parents=True, exist_ok=True)
+
+    cleanup_old_sessions(base_path)
+
+    clone_url = f"https://x-access-token:{github_token}@github.com/{repository}.git"
+
+    if cache_path.exists() and (cache_path / ".git").exists():
+        logger.info(f"Found cached repo at {cache_path}, copying to {session_path}")
+        try:
+            if session_path.exists():
+                shutil.rmtree(session_path)
+            shutil.copytree(cache_path, session_path)
+
+            subprocess.run(
+                ["git", "fetch", "origin"],
+                cwd=str(session_path),
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            subprocess.run(
+                ["git", "reset", "--hard", "origin/HEAD"],
+                cwd=str(session_path),
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            logger.info(f"Successfully reused cached repo and updated to latest")
+        except Exception as e:
+            logger.warning(f"Failed to reuse cache: {e}, will clone fresh")
+            if session_path.exists():
+                shutil.rmtree(session_path)
+            _clone_fresh(clone_url, session_path, repository)
+    else:
+        logger.info(f"No cache found, cloning fresh into {session_path}")
+        _clone_fresh(clone_url, session_path, repository)
 
-    if repo_path.exists():
-        logger.info(f"Directory {repo_path} already exists, removing it (Lambda container reuse)")
         try:
-            shutil.rmtree(repo_path)
-            logger.info(f"Successfully removed existing directory at {repo_path}")
+            if cache_path.exists():
+                shutil.rmtree(cache_path)
+            shutil.copytree(session_path, cache_path)
+            logger.info(f"Cached repo for future use at {cache_path}")
         except Exception as e:
-            logger.error(f"Failed to remove existing directory at {repo_path}: {e}")
-            raise RuntimeError(f"Failed to clean up existing directory: {e}")
+            logger.warning(f"Failed to create cache: {e}, continuing without cache")
 
-    clone_url = f"https://x-access-token:{github_token}@github.com/{repository}.git"
+    configure_git_auth(str(session_path))
+
+    return str(session_path)
 
-    logger.info(f"Cloning {repository} into {repo_path}")
+
+def _clone_fresh(clone_url: str, repo_path: Path, repository: str) -> None:
+    """Helper to clone a repository from scratch."""
     try:
         subprocess.run(
             ["git", "clone", clone_url, str(repo_path)],
             check=True,
             capture_output=True,
             text=True,
         )
+        logger.info(f"Successfully cloned {repository}")
     except subprocess.CalledProcessError as e:
         logger.error(f"Failed to clone repository: {e.stderr}")
         raise RuntimeError(f"Failed to clone {repository}: {e.stderr}")
 
-    configure_git_auth(str(repo_path))
-
-    return str(repo_path)
-
 
 def configure_git_auth(repo_path: str) -> None:
     """Configure git user for the repository."""