From 2e1d16ef9ebe60261ba76cf63e8bd3296e4f5924 Mon Sep 17 00:00:00 2001
From: patrick <patrickjaroszewski@protonmail.com>
Date: Mon, 9 Mar 2026 14:25:24 -0400
Subject: [PATCH] refactor: move db and file watcher setup to use async context
 manager

---
 backend/src/chitai/app.py | 41 ++++++++++++++++++++++++++-------------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/backend/src/chitai/app.py b/backend/src/chitai/app.py
index d82b0c7..102a957 100644
--- a/backend/src/chitai/app.py
+++ b/backend/src/chitai/app.py
@@ -1,5 +1,6 @@
 import asyncio
-from typing import Any
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator
 
 from chitai.services.book import BookService
 from chitai.services.consume import ConsumeDirectoryWatcher
@@ -63,14 +64,15 @@ oauth2_auth = OAuth2PasswordBearerAuth[User](
         "/access/signup",
         "/opds",
         "/schema",
+        "/syncs",
+        "/users/auth",
     ],
 )
 
 watcher_task: asyncio.Task
 
-async def startup():
-    """Run setup."""
-
+@asynccontextmanager
+async def setup_db_connection(app: Litestar) -> AsyncGenerator[None, None]:
     # Setup databse
     async with settings.alchemy_config.get_session() as db_session:
         # Create default library if none exist
@@ -86,22 +88,31 @@ async def startup():
             )
             await db_session.commit()
 
-        # book_service = BookService(session=db_session)
+        try:
+            yield
+        finally:
+            await db_session.aclose()
 
+
+@asynccontextmanager
+async def setup_directory_watcher(app: Litestar) -> AsyncGenerator[None, None]:
+    
     # Create book covers directory if it does not exist
     await create_directory(settings.book_cover_path)
-
     # Create consume directory
     await create_directory(settings.consume_path)
 
+    async with settings.alchemy_config.get_session() as db_session:
+        book_service = BookService(session=db_session)
+        library_service = LibraryService(session=db_session)
 
-    # file_watcher = ConsumeDirectoryWatcher(settings.consume_path, library_service, book_service)
-    # watcher_task = asyncio.create_task(file_watcher.init_watcher())
-
-async def shutdown():
-    """ Run shutdown tasks. """
-
-    watcher_task.cancel()
+        file_watcher = ConsumeDirectoryWatcher(settings.consume_path, library_service, book_service)
+        watcher_task = asyncio.create_task(file_watcher.init_watcher())
+        
+        try:
+            yield
+        finally:
+            watcher_task.cancel()
 
 def create_app() -> Litestar:
     return Litestar(
@@ -114,12 +125,14 @@ def create_app() -> Litestar:
             c.PublisherController,
             c.TagController,
             c.OpdsController,
+            c.DeviceController,
+            c.KosyncController,
             create_static_files_router(path="/covers", directories=["./covers"]),
             index,
             healthcheck,
         ],
         exception_handlers=exception_handlers,
-        on_startup=[startup],
+        lifespan=[setup_db_connection, setup_directory_watcher],
         plugins=[alchemy],
         on_app_init=[oauth2_auth.on_app_init],
         openapi_config=OpenAPIConfig(