refactor(controllers): adopt with_session in web saved_message (#41859)

api/controllers/web/saved_message.py

@@ -1,15 +1,16 @@
 from uuid import UUID
 
 from pydantic import TypeAdapter
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import NotFound
 
 from controllers.common.controller_schemas import SavedMessageCreatePayload, SavedMessageListQuery
 from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
+from controllers.console.app.wraps import with_session
 from controllers.console.wraps import model_validate
 from controllers.web import web_ns
 from controllers.web.error import NotCompletionAppError
 from controllers.web.wraps import WebApiResource
-from extensions.ext_database import db
 from fields.conversation_fields import MessageResponseSource, ResultResponse
 from fields.message_fields import SavedMessageInfiniteScrollPagination, SavedMessageItem
 from models.model import App, EndUser
@@ -36,12 +37,12 @@ class SavedMessageListApi(WebApiResource):
         }
     )
     @web_ns.response(200, "Success", web_ns.models[SavedMessageInfiniteScrollPagination.__name__])
+    @with_session(write=False)
     @model_validate(SavedMessageListQuery)
-    def get(self, query: SavedMessageListQuery, app_model: App, end_user: EndUser):
+    def get(self, query: SavedMessageListQuery, session: Session, app_model: App, end_user: EndUser):
         if app_model.mode != "completion":
             raise NotCompletionAppError()
 
-        session = db.session()
         pagination = SavedMessageService.pagination_by_last_id(
             app_model, end_user, query.last_id, query.limit, session=session
         )
@@ -73,13 +74,14 @@ class SavedMessageListApi(WebApiResource):
     )
     @web_ns.response(200, "Message saved successfully", web_ns.models[ResultResponse.__name__])
     @web_ns.expect(web_ns.models[SavedMessageCreatePayload.__name__])
+    @with_session
     @model_validate(SavedMessageCreatePayload)
-    def post(self, payload: SavedMessageCreatePayload, app_model: App, end_user: EndUser):
+    def post(self, payload: SavedMessageCreatePayload, session: Session, app_model: App, end_user: EndUser):
         if app_model.mode != "completion":
             raise NotCompletionAppError()
 
         try:
-            SavedMessageService.save(app_model, end_user, payload.message_id, session=db.session())
+            SavedMessageService.save(app_model, end_user, payload.message_id, session=session)
         except MessageNotExistsError:
             raise NotFound("Message Not Exists.")
 
@@ -102,12 +104,13 @@ class SavedMessageApi(WebApiResource):
         }
     )
     @web_ns.response(204, "Message removed successfully")
-    def delete(self, app_model: App, end_user: EndUser, message_id: UUID):
+    @with_session
+    def delete(self, session: Session, app_model: App, end_user: EndUser, message_id: UUID):
         message_id_str = str(message_id)
 
         if app_model.mode != "completion":
             raise NotCompletionAppError()
 
-        SavedMessageService.delete(app_model, end_user, message_id_str, session=db.session())
+        SavedMessageService.delete(app_model, end_user, message_id_str, session=session)
 
         return "", 204

api/tests/unit_tests/controllers/web/test_saved_message.py

@@ -36,9 +36,9 @@ def _end_user() -> EndUser:
     )
 
 
-# The @model_validate decorator wraps the handler; tests call the undecorated
-# function so they can pass the validated pydantic payload directly and skip
-# the flask request-parsing step.
+# The @with_session and @model_validate decorators wrap the handler; tests call
+# the undecorated function so they can pass the validated pydantic payload and a
+# stub session directly and skip the flask request-parsing step.
 _list_get = inspect.unwrap(SavedMessageListApi.get)
 _list_post = inspect.unwrap(SavedMessageListApi.post)
 
@@ -51,7 +51,7 @@ class TestSavedMessageListApiGet:
         query = SavedMessageListQuery.model_validate({})
         with app.test_request_context("/saved-messages"):
             with pytest.raises(NotCompletionAppError):
-                _list_get(SavedMessageListApi(), query, _chat_app(), _end_user())
+                _list_get(SavedMessageListApi(), query, MagicMock(), _chat_app(), _end_user())
 
     @patch("controllers.web.saved_message.SavedMessageService.pagination_by_last_id")
     def test_happy_path(self, mock_paginate: MagicMock, app: Flask) -> None:
@@ -59,7 +59,7 @@ class TestSavedMessageListApiGet:
         query = SavedMessageListQuery.model_validate({"limit": 20})
 
         with app.test_request_context("/saved-messages?limit=20"):
-            result = _list_get(SavedMessageListApi(), query, _completion_app(), _end_user())
+            result = _list_get(SavedMessageListApi(), query, MagicMock(), _completion_app(), _end_user())
 
         assert result["limit"] == 20
         assert result["has_more"] is False
@@ -73,14 +73,14 @@ class TestSavedMessageListApiPost:
         payload = SavedMessageCreatePayload.model_validate({"message_id": str(uuid4())})
         with app.test_request_context("/saved-messages", method="POST"):
             with pytest.raises(NotCompletionAppError):
-                _list_post(SavedMessageListApi(), payload, _chat_app(), _end_user())
+                _list_post(SavedMessageListApi(), payload, MagicMock(), _chat_app(), _end_user())
 
     @patch("controllers.web.saved_message.SavedMessageService.save")
     def test_save_success(self, mock_save: MagicMock, app: Flask) -> None:
         payload = SavedMessageCreatePayload.model_validate({"message_id": str(uuid4())})
 
         with app.test_request_context("/saved-messages", method="POST"):
-            result = _list_post(SavedMessageListApi(), payload, _completion_app(), _end_user())
+            result = _list_post(SavedMessageListApi(), payload, MagicMock(), _completion_app(), _end_user())
 
         assert result["result"] == "success"
 
@@ -90,7 +90,7 @@ class TestSavedMessageListApiPost:
 
         with app.test_request_context("/saved-messages", method="POST"):
             with pytest.raises(NotFound, match="Message Not Exists"):
-                _list_post(SavedMessageListApi(), payload, _completion_app(), _end_user())
+                _list_post(SavedMessageListApi(), payload, MagicMock(), _completion_app(), _end_user())
 
 
 # ---------------------------------------------------------------------------