fix: stabilize logging and RabbitMQ shutdown lifecycle

src/memos/api/lifecycle.py

@@ -0,0 +1,26 @@
+from collections.abc import Mapping
+from typing import Any
+
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def shutdown_components(components: Mapping[str, Any] | None) -> None:
+    """Release long-lived API components before the logging system shuts down."""
+    if not components:
+        return
+
+    mem_scheduler = components.get("mem_scheduler")
+    if mem_scheduler is None:
+        return
+
+    for method_name in ("stop", "rabbitmq_close"):
+        method = getattr(mem_scheduler, method_name, None)
+        if not callable(method):
+            continue
+        try:
+            method()
+        except Exception:
+            logger.exception("Failed to run mem_scheduler.%s during API shutdown", method_name)

src/memos/api/server_api.py

@@ -1,14 +1,18 @@
 import logging
 import os
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
 from dotenv import load_dotenv
 from fastapi import FastAPI, HTTPException
 from fastapi.exceptions import RequestValidationError
 from starlette.staticfiles import StaticFiles
 
 from memos.api.exceptions import APIExceptionHandler
+from memos.api.lifecycle import shutdown_components
 from memos.api.middleware.request_context import RequestContextMiddleware
-from memos.api.routers.server_router import router as server_router
+from memos.api.routers import server_router as server_router_module
 from memos.plugins.manager import plugin_manager
 
 
@@ -25,17 +29,25 @@
     os.getenv("MEMSCHEDULER_REDIS_STREAM_KEY_PREFIX"),
 )
 
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
+    yield
+    shutdown_components(server_router_module.components)
+
+
 app = FastAPI(
     title="MemOS Server REST APIs",
     description="A REST API for managing multiple users with MemOS Server.",
     version="1.0.1",
+    lifespan=lifespan,
 )
 
 app.mount("/download", StaticFiles(directory=os.getenv("FILE_LOCAL_PATH")), name="static_mapping")
 
 app.add_middleware(RequestContextMiddleware, source="server_api")
 # Include routers
-app.include_router(server_router)
+app.include_router(server_router_module.router)
 
 
 @app.get("/health")

src/memos/api/server_api_ext.py

@@ -18,6 +18,9 @@
 import logging
 import os
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI
 from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
@@ -26,11 +29,12 @@
 from starlette.responses import Response
 
 # Import Krolik extensions
+from memos.api.lifecycle import shutdown_components
 from memos.api.middleware.rate_limit import RateLimitMiddleware
-from memos.api.routers.admin_router import router as admin_router
 
 # Import base routers from MemOS
-from memos.api.routers.server_router import router as server_router
+from memos.api.routers import server_router as server_router_module
+from memos.api.routers.admin_router import router as admin_router
 
 
 # Try to import exception handlers (may vary between MemOS versions)
@@ -58,11 +62,18 @@ async def dispatch(self, request: Request, call_next) -> Response:
         return response
 
 
+@asynccontextmanager
+async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
+    yield
+    shutdown_components(server_router_module.components)
+
+
 # Create FastAPI app
 app = FastAPI(
     title="MemOS Server REST APIs (Krolik Extended)",
     description="MemOS API with authentication, rate limiting, and admin endpoints.",
     version="2.0.3-krolik",
+    lifespan=lifespan,
 )
 
 # CORS configuration
@@ -94,7 +105,7 @@ async def dispatch(self, request: Request, call_next) -> Response:
     logger.info("Rate limiting enabled")
 
 # Include routers
-app.include_router(server_router)
+app.include_router(server_router_module.router)
 app.include_router(admin_router)
 
 # Exception handlers

src/memos/log.py

@@ -27,6 +27,8 @@
 load_dotenv()
 
 selected_log_level = logging.DEBUG if settings.DEBUG else logging.WARNING
+_LOGGING_CONFIG_LOCK = threading.RLock()
+_LOGGING_CONFIGURED_PID: int | None = None
 
 
 def _setup_logfile() -> Path:
@@ -224,12 +226,31 @@ def close(self):
 }
 
 
+def _get_current_pid() -> int:
+    return os.getpid()
+
+
+def configure_logging(force: bool = False) -> None:
+    """Configure process-local logging once.
+
+    Re-running dictConfig replaces and closes existing handlers. Guarding it avoids races with
+    background threads that may be emitting log records while other modules import loggers.
+    """
+    global _LOGGING_CONFIGURED_PID
+
+    with _LOGGING_CONFIG_LOCK:
+        current_pid = _get_current_pid()
+        if force or current_pid != _LOGGING_CONFIGURED_PID:
+            dictConfig(LOGGING_CONFIG)
+            _LOGGING_CONFIGURED_PID = current_pid
+
+
 def get_logger(name: str | None = None) -> logging.Logger:
     """returns the project logger, scoped to a child name if provided
     Args:
         name: will define a child logger
     """
-    dictConfig(LOGGING_CONFIG)
+    configure_logging()
 
     parent_logger = logging.getLogger("")
     if name:

tests/api/test_lifecycle.py

@@ -0,0 +1,38 @@
+from memos.api.lifecycle import shutdown_components
+
+
+class SchedulerStub:
+    def __init__(self, fail_stop: bool = False):
+        self.fail_stop = fail_stop
+        self.calls = []
+
+    def stop(self):
+        self.calls.append("stop")
+        if self.fail_stop:
+            raise RuntimeError("stop failed")
+
+    def rabbitmq_close(self):
+        self.calls.append("rabbitmq_close")
+
+
+def test_shutdown_components_stops_scheduler_before_rabbitmq_close():
+    scheduler = SchedulerStub()
+
+    shutdown_components({"mem_scheduler": scheduler})
+
+    assert scheduler.calls == ["stop", "rabbitmq_close"]
+
+
+def test_shutdown_components_still_closes_rabbitmq_when_stop_fails():
+    scheduler = SchedulerStub(fail_stop=True)
+
+    shutdown_components({"mem_scheduler": scheduler})
+
+    assert scheduler.calls == ["stop", "rabbitmq_close"]
+
+
+def test_shutdown_components_skips_missing_methods():
+    class MinimalScheduler:
+        pass
+
+    shutdown_components({"mem_scheduler": MinimalScheduler()})

tests/test_log.py

@@ -28,3 +28,34 @@ def test_get_logger_returns_logger():
     assert any(isinstance(h, logging.StreamHandler) for h in logger.parent.handlers) or any(
         isinstance(h, logging.FileHandler) for h in logger.parent.handlers
     )
+
+
+def test_get_logger_configures_logging_once_per_process(monkeypatch):
+    calls = []
+
+    monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None)
+    monkeypatch.setattr(log, "_get_current_pid", lambda: 123)
+    monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config))
+
+    log.get_logger("first")
+    log.get_logger("second")
+
+    assert len(calls) == 1
+
+
+def test_get_logger_reconfigures_after_process_fork(monkeypatch):
+    calls = []
+    pid = 123
+
+    def getpid():
+        return pid
+
+    monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None)
+    monkeypatch.setattr(log, "_get_current_pid", getpid)
+    monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config))
+
+    log.get_logger("parent")
+    pid = 456
+    log.get_logger("child")
+
+    assert len(calls) == 2