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:
@@ -27,10 +27,14 @@ def compress_imgfile(file, max_size):
|
||||
img = Image.open(file)
|
||||
rgb_image = img.convert("RGB")
|
||||
quality = 95
|
||||
min_quality = 10
|
||||
while True:
|
||||
out_buf = io.BytesIO()
|
||||
rgb_image.save(out_buf, "JPEG", quality=quality)
|
||||
if fsize(out_buf) <= max_size:
|
||||
if fsize(out_buf) <= max_size or quality <= min_quality:
|
||||
# Stop at min_quality: further decrements would pass an invalid
|
||||
# quality (<1) to PIL and the loop would otherwise never terminate
|
||||
# for images that cannot be compressed below max_size.
|
||||
return out_buf
|
||||
quality -= 5
|
||||
|
||||
|
||||
Reference in New Issue
Block a user