From ba08ca2d2e58b03db4b40dec69e13096985ef28f Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Thu, 29 Aug 2024 19:18:33 +0100
Subject: [PATCH] Changed to memory logger

---
 api_server/routes/internal/internal_routes.py | 10 ++------
 app/log_interceptor.py                        | 21 ----------------
 app/logger.py                                 | 25 +++++++++++++++++++
 main.py                                       |  7 ++----
 4 files changed, 29 insertions(+), 34 deletions(-)
 delete mode 100644 app/log_interceptor.py
 create mode 100644 app/logger.py

diff --git a/api_server/routes/internal/internal_routes.py b/api_server/routes/internal/internal_routes.py
index c998f1e0a..8c46215f0 100644
--- a/api_server/routes/internal/internal_routes.py
+++ b/api_server/routes/internal/internal_routes.py
@@ -2,7 +2,7 @@ from aiohttp import web
 from typing import Optional
 from folder_paths import models_dir, user_directory, output_directory
 from api_server.services.file_service import FileService
-from app.log_interceptor import logs
+import app.logger
 
 class InternalRoutes:
     '''
@@ -34,13 +34,7 @@ class InternalRoutes:
 
         @self.routes.get('/logs')
         async def get_logs(request):
-            return web.json_response([
-                {
-                    'timestamp': dt.isoformat(),
-                    'message': message
-                }
-                for dt, message in logs
-            ])
+            return web.json_response(app.logger.get_logs())
 
     def get_app(self):
         if self._app is None:
diff --git a/app/log_interceptor.py b/app/log_interceptor.py
deleted file mode 100644
index 011aac701..000000000
--- a/app/log_interceptor.py
+++ /dev/null
@@ -1,21 +0,0 @@
-import io
-import threading
-from datetime import datetime
-
-from collections import deque
-logs = deque(maxlen=300)
-
-class LogInterceptor(io.TextIOWrapper):
-    def __init__(self, stream,  *args, **kwargs):
-        buffer = stream.buffer
-        encoding = stream.encoding
-        super().__init__(buffer, *args, **kwargs, encoding=encoding)
-        self._lock = threading.Lock()
-
-    def write(self, data):
-        with self._lock:
-            logs.append((datetime.now(), data))
-        super().write(data)
-
-    def flush(self):
-        super().flush()
diff --git a/app/logger.py b/app/logger.py
new file mode 100644
index 000000000..a1287c999
--- /dev/null
+++ b/app/logger.py
@@ -0,0 +1,25 @@
+import logging
+from logging.handlers import MemoryHandler
+from collections import deque
+
+logs = None
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+
+def get_logs():
+    return '\n'.join([formatter.format(x) for x in logs])
+
+def setup_logger(capacity=300):
+    global logs
+    if logs: return
+
+    # Setup default global logger
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    logger.addHandler(logging.StreamHandler())
+
+    # Create a memory handler with a deque as its buffer
+    logs = deque(maxlen=capacity)
+    memory_handler = MemoryHandler(capacity, flushLevel=logging.INFO)
+    memory_handler.buffer = logs
+    memory_handler.setFormatter(formatter)
+    logger.addHandler(memory_handler)
diff --git a/main.py b/main.py
index 8c5d64bbe..c75c0e75e 100644
--- a/main.py
+++ b/main.py
@@ -1,11 +1,8 @@
 import comfy.options
 comfy.options.enable_args_parsing()
 
-# Intercept logs for error reporting
-import sys
-from app.log_interceptor import LogInterceptor
-sys.stdout = LogInterceptor(sys.stdout)
-sys.stderr = LogInterceptor(sys.stderr)
+from app.logger import setup_logger
+setup_logger()
 
 import os
 import importlib.util