fix(mcp): scan all message DB shards in get_chat_images (#84)

WeChat rolls a chat's messages over to the next `message_N.db` shard
once the current shard fills up (~100 MB), so any chat older than the
current shard window has its history split across multiple shards. The
other message-query tools — `get_chat_history`, `search_messages`, and
`decode_image` — already iterate every matching shard via the plural
helper `_find_msg_tables_for_user`. Only `get_chat_images` still used
the singular `_find_msg_table_for_user`, which returns the first shard
that contains the user's table.

Effect: every image that lived in a non-first shard was silently
dropped from `get_chat_images`. On a long-lived chat with many images,
the tool would return only the most recent slice and pretend the rest
did not exist.

Fix: switch `get_chat_images` to `_find_msg_tables_for_user`, fetch
`limit` images per shard, merge, sort by `create_time` DESC, and slice
to `limit`. This mirrors how the other tools fan out across shards.

Tests in `tests/test_get_chat_images_multishard.py`:

- `test_collects_images_from_every_shard` — both shards' images appear
  in the output (the regression case)
- `test_global_sort_by_create_time_desc` — newer image from an older
  shard still wins, output is globally sorted (not per-shard concat)
- `test_limit_truncates_globally_across_shards` — limit=3 takes the 3
  newest overall, not "first shard wins"
- `test_no_shards_returns_not_found` — empty shard list path
- `test_all_shards_empty_returns_no_images` — every shard empty path

All 156 tests pass locally (151 baseline + 5 new). Public tool
signature is unchanged; only the internal scanning loop is widened.
This commit is contained in:
Belugary
2026-05-12 16:18:49 +08:00
committed by GitHub
parent c162a9b92f
commit cd329afa1b
2 changed files with 151 additions and 4 deletions

View File

@@ -2524,14 +2524,28 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str:
names = get_contact_names()
display_name = names.get(username, username)
db_path, table_name = _find_msg_table_for_user(username)
if not db_path:
# 同 chat 的消息会分散在多个 message_N.db shard 里 (上限 ~100MB/shard 时滚动到下一个);
# 单 shard 查找会漏掉其他 shard 的图片。其他工具 (get_chat_history / search_messages /
# decode_image) 早已用复数版本 scan 全部 shard, 这里对齐一致。
shards = _find_msg_tables_for_user(username)
if not shards:
return f"找不到 {display_name} 的消息记录"
images = _image_resolver.list_chat_images(db_path, table_name, username, limit)
if not images:
# 每个 shard 取 limit 张, 合并后按 create_time DESC 全局排序, 取最新 limit 张。
# 单 shard 至少够本次返回, 避免一个 shard 凑不出 limit 时其他 shard 没机会贡献。
all_images = []
for shard in shards:
shard_images = _image_resolver.list_chat_images(
shard['db_path'], shard['table_name'], username, limit
)
all_images.extend(shard_images)
if not all_images:
return f"{display_name} 无图片消息"
all_images.sort(key=lambda img: img['create_time'], reverse=True)
images = all_images[:limit]
lines = []
for img in images:
time_str = datetime.fromtimestamp(img['create_time']).strftime('%Y-%m-%d %H:%M')