From b61e9c9177df14113fc7592a7a70b7e7da1f6324 Mon Sep 17 00:00:00 2001
From: UladzislauK-Writer <uladzislau.kasabutski@writer.com>
Date: Tue, 3 Feb 2026 13:10:20 +0100
Subject: [PATCH] fix: make custom server setup more resilient

---
 src/writer/app_runner.py |  3 +++
 src/writer/serve.py      | 16 +++++++++++++++-
 src/writer/ss_types.py   |  8 ++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/writer/app_runner.py b/src/writer/app_runner.py
index 115140880..5395a70c3 100644
--- a/src/writer/app_runner.py
+++ b/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()
diff --git a/src/writer/serve.py b/src/writer/serve.py
index 0dc02a22c..5f00838bb 100644
--- a/src/writer/serve.py
+++ b/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
 
diff --git a/src/writer/ss_types.py b/src/writer/ss_types.py
index 68e83feeb..43b122710 100644
--- a/src/writer/ss_types.py
+++ b/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):