fix: catch oom

server/repl.py

@@ -32,6 +32,12 @@
 log_lock = asyncio.Lock()
 
 
+class ReplOOMError(ReplError):
+    """Raised when the REPL exceeds its memory limit."""
+
+    pass
+
+
 async def log_snippet(uuid: UUID, snippet_id: str, code: str) -> None:
     if settings.environment == Environment.prod:
         header = f"[{uuid.hex[:8]}] Running snippet {snippet_id}:"
@@ -75,7 +81,9 @@ def __init__(
         self.header_cmd_response: ReplResponse | None = None
 
         self.proc: Process | None = None
-        self.error_file = tempfile.TemporaryFile("w+")
+        self.error_file = tempfile.TemporaryFile(
+            "w+", encoding="utf-8", errors="ignore"
+        )
         self.max_memory_bytes = max_repl_mem * 1024 * 1024
         self.max_repl_uses = max_repl_uses
 
@@ -146,6 +154,9 @@ def _preexec() -> None:
 
             os.setsid()
 
+        self.error_file.seek(0)
+        self.error_file.truncate()
+
         self.proc = await asyncio.create_subprocess_exec(
             "lake",
             "env",
@@ -154,7 +165,7 @@ def _preexec() -> None:
             env=os.environ,
             stdin=asyncio.subprocess.PIPE,
             stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE,
+            stderr=self.error_file,
             preexec_fn=_preexec,
         )
 
@@ -231,10 +242,13 @@ async def send_timeout(
             )
             raise e
         except LeanError as e:
-            logger.exception("Lean REPL error: %s", e)
+            logger.exception("Lean REPL error: {}", e)
             raise e
+        except ReplOOMError as e:
+            logger.error("[{}] REPL OOM: {}", self.uuid.hex[:8], e)
+            return ReplResponse(id=snippet.id, error=str(e), time=0.0)
         except ReplError as e:
-            logger.exception("REPL error: %s", e)
+            logger.exception("REPL error: {}", e)
             raise e
 
         return ReplResponse(
@@ -287,26 +301,34 @@ async def send(
             logger.error("Broken pipe while writing to REPL stdin")
             raise LeanError("Lean process broken pipe")
         except Exception as e:
-            logger.error("Failed to write to REPL stdin: %s", e)
+            logger.error("Failed to write to REPL stdin: {}", e)
             raise LeanError("Failed to write to REPL stdin")
 
         logger.debug("Reading response from REPL stdout")
         raw = await self._read_response()
         elapsed = loop.time() - start
 
-        logger.debug("Raw response from REPL: %r", raw)
+        if self.proc.returncode not in (None, 0):
+            self.error_file.seek(0)
+            err = (
+                self.error_file.read().strip()
+                or "Lean REPL died, probably out of memory"
+            )
+            raise ReplOOMError(err)
+
+        logger.debug("Raw response from REPL: {}", raw)
         try:
             resp: CommandResponse | Error = json.loads(raw)
         except json.JSONDecodeError:
-            logger.error("JSON decode error: %r", raw)
+            logger.error("JSON decode error: {}", raw)
             raise ReplError("JSON decode error")
 
         self.error_file.seek(0)
         err = self.error_file.read().strip()
         self.error_file.seek(0)
         self.error_file.truncate(0)
         if err:
-            logger.error("Stderr: %s", err)
+            logger.error("Stderr: {}", err)
             raise LeanError(err)
 
         elapsed_time = round(elapsed, 6)
@@ -336,7 +358,7 @@ async def _read_response(self) -> bytes:
                     break
                 lines.append(chunk)
         except Exception as e:
-            logger.error("Failed to read from REPL stdout: %s", e)
+            logger.error("Failed to read from REPL stdout: {}", e)
             raise LeanError("Failed to read from REPL stdout")
         return b"".join(lines)