mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-07-17 11:07:11 +08:00
fix: avoid KeyError on cancel and infinite loop in image compression
Two independent robustness fixes: 1. ChatChannel.cancel_session / cancel_all_session raised KeyError when a session existed in self.sessions but no future had been dispatched yet. self.sessions[sid] is created in produce(), but self.futures[sid] is only created later in consume() on first dispatch. Cancelling in that window (e.g. user sends a message then immediately cancels) crashed the cancel path. Use self.futures.get(sid, []) so an absent entry is a no-op. 2. compress_imgfile decremented JPEG quality by 5 with no lower bound. For an image that cannot be compressed below max_size, quality went 0, negative, ... — the loop never terminated and passed invalid quality values to PIL. Add a min_quality floor (10) and return the best effort once reached. Adds tests/test_robustness_fixes.py covering both paths (5 tests).
This commit is contained in:
@@ -519,7 +519,10 @@ class ChatChannel(Channel):
|
||||
def cancel_session(self, session_id):
|
||||
with self.lock:
|
||||
if session_id in self.sessions:
|
||||
for future in self.futures[session_id]:
|
||||
# futures[session_id] is only created in consume() when a task is
|
||||
# dispatched, so it may be absent if cancel happens right after
|
||||
# produce() but before the first dispatch. Default to [].
|
||||
for future in self.futures.get(session_id, []):
|
||||
future.cancel()
|
||||
cnt = self.sessions[session_id][0].qsize()
|
||||
if cnt > 0:
|
||||
@@ -529,7 +532,7 @@ class ChatChannel(Channel):
|
||||
def cancel_all_session(self):
|
||||
with self.lock:
|
||||
for session_id in self.sessions:
|
||||
for future in self.futures[session_id]:
|
||||
for future in self.futures.get(session_id, []):
|
||||
future.cancel()
|
||||
cnt = self.sessions[session_id][0].qsize()
|
||||
if cnt > 0:
|
||||
|
||||
Reference in New Issue
Block a user