Refactor message handling: improve session management and add timeout configuration

Author: Keniis0712

StealthIM/apis/__init__.py

@@ -2,4 +2,4 @@
 from . import user
 from . import util
 from . import group
-from . import message
+from . import message

StealthIM/apis/message.py

@@ -8,6 +8,12 @@
 from .common import request, NoValResult, Result
 from .. import logger
 
+timeout_for_sse = aiohttp.ClientTimeout(
+    total=None,  # 整体无超时
+    sock_connect=None,  # 连接无超时
+    sock_read=None  # 读取无超时
+)
+
 
 class MessageType(enum.Enum):
     Text = 0
@@ -105,29 +111,31 @@ async def get_message(
         "Authorization": f"Bearer {session}",
     }
 
-    async with aiohttp.ClientSession(headers=headers) as session:
-        async with session.get(api_address) as response:
-            if response.status != 200:
-                raise RuntimeError(f"Request failed with status: {response.status}")
-
-            async for line in response.content:
-                # 解码成字符串
-                line = line.decode('utf-8').strip()
-                if line.startswith('data:'):
-                    # 提取消息内容
-                    message = line[len('data:'):].strip()
-
-                    data = json.loads(message)
-                    logger.debug(f"Response data: {data}")
-                    if data["result"]["code"] != 800:
-                        raise RuntimeError(f"Request failed with code: {data['result']['code']}")
-                    for msg in data["msg"]:
-                        yield Message(
-                            groupid=msg["groupid"],
-                            msg=msg["msg"],
-                            msgid=msg["msgid"],
-                            time=msg["time"],
-                            type=MessageType(msg["type"]),
-                            username=msg["username"],
-                            hash=msg.get("hash", None),
-                        )
+    async with (
+        aiohttp.ClientSession(headers=headers, timeout=timeout_for_sse) as session,
+        session.get(api_address) as response
+    ):
+        if response.status != 200:
+            raise RuntimeError(f"Request failed with status: {response.status}")
+
+        async for line in response.content:
+            # 解码成字符串
+            line = line.decode('utf-8').strip()
+            if line.startswith('data:'):
+                # 提取消息内容
+                message = line[len('data:'):].strip()
+
+                data = json.loads(message)
+                logger.debug(f"Response data: {data}")
+                if data["result"]["code"] != 800:
+                    raise RuntimeError(f"Request failed with code: {data['result']['code']}")
+                for msg in data["msg"]:
+                    yield Message(
+                        groupid=msg["groupid"],
+                        msg=msg["msg"],
+                        msgid=msg["msgid"],
+                        time=msg["time"],
+                        type=MessageType(msg["type"]),
+                        username=msg["username"],
+                        hash=msg.get("hash", None),
+                    )

StealthIM/group.py

@@ -33,7 +33,6 @@ async def join(self, password: str) -> JoinGroupResult:
         Join a Group.
 
         Args:
-            user (User): The user to join.
             password (str): The password.
 
         Returns:
@@ -44,7 +43,6 @@ async def join(self, password: str) -> JoinGroupResult:
         """
         return await StealthIM.apis.group.join_group(self.user.server.url, self.user.session, self.group_id, password)
 
-
     async def get_members(self) -> GroupInfoResult:
         """
         Get members of the group.