fix: make custom server setup more resilient (#1277)

Author: UladzislauK-Writer

src/writer/app_runner.py

@@ -171,6 +171,9 @@ def _handle_session_init(
         if session is None:
             raise MessageHandlingException("Session rejected.")
 
+        for mail in payload.additionalMail:
+            session.session_state.add_log_entry(**mail.model_dump())
+
         user_state = {}
         try:
             user_state = session.session_state.user_state.to_dict()

src/writer/serve.py

@@ -13,6 +13,7 @@
 import tempfile
 import textwrap
 import time
+import traceback
 import typing
 from contextlib import asynccontextmanager, suppress
 from importlib.machinery import ModuleSpec
@@ -121,6 +122,8 @@ def get_asgi_app(
     global app
     if serve_mode not in ["run", "edit"]:
         raise ValueError("""Invalid mode. Must be either "run" or "edit".""")
+    
+    custom_server_setup_mail: list[dict] = []
 
     _fix_mimetype()
     app_runner = AppRunner(user_app_path, serve_mode)
@@ -351,6 +354,7 @@ async def init(
                 cookies=dict(request.cookies),
                 headers=dict(request.headers),
                 proposedSessionId=initBody.proposedSessionId,
+                additionalMail=custom_server_setup_mail,
             )
         )
 
@@ -960,7 +964,17 @@ async def stream(websocket: WebSocket):
 
     # Return
     if enable_server_setup is True:
-        _execute_server_setup_hook(user_app_path)
+        try:
+            _execute_server_setup_hook(user_app_path)
+        except Exception as e:
+            custom_server_setup_mail.append(
+                {
+                    "type": "error",
+                    "title": "Custom server setup error",
+                    "message": str(e),
+                    "code": traceback.format_exc()
+                }
+            )
 
     return app
 

src/writer/ss_types.py

@@ -139,10 +139,18 @@ class AppProcessServerRequest(BaseModel):
     payload: Optional[Any] = None
 
 
+class AdditionalMail(BaseModel):
+    type: Literal["info", "error"]
+    title: str
+    message: str
+    code: Optional[str] = None
+
+
 class InitSessionRequestPayload(BaseModel):
     cookies: Optional[Dict[str, str]] = None
     headers: Optional[Dict[str, str]] = None
     proposedSessionId: Optional[str] = None
+    additionalMail: list[AdditionalMail] = []
 
 
 class InitSessionRequest(AppProcessServerRequest):