diff --git a/decode_image.py b/decode_image.py index 43e8e89..e7cb0fa 100644 --- a/decode_image.py +++ b/decode_image.py @@ -450,17 +450,30 @@ class ImageResolver: 'size': os.path.getsize(final_path), } - def list_chat_images(self, db_path, table_name, username, limit=20): - """列出某个聊天中的所有图片消息""" + def list_chat_images(self, db_path, table_name, username, limit=20, start_ts=None, end_ts=None): + """列出某个聊天中的所有图片消息 + + 可选 start_ts / end_ts (unix 秒) 过滤时间范围。 + """ + clauses = ['local_type = 3'] + params = [] + if start_ts is not None: + clauses.append('create_time >= ?') + params.append(start_ts) + if end_ts is not None: + clauses.append('create_time <= ?') + params.append(end_ts) + params.append(limit) + where_sql = ' AND '.join(clauses) conn = sqlite3.connect(db_path) try: rows = conn.execute(f""" SELECT local_id, create_time FROM [{table_name}] - WHERE local_type = 3 + WHERE {where_sql} ORDER BY create_time DESC LIMIT ? - """, (limit,)).fetchall() + """, params).fetchall() except Exception as e: conn.close() return [] diff --git a/mcp_server.py b/mcp_server.py index 5c7ea3a..8c6ca45 100644 --- a/mcp_server.py +++ b/mcp_server.py @@ -2830,7 +2830,7 @@ def decode_transfer(chat_name: str, local_id: int, create_time: int = 0) -> str: @mcp.tool() -def get_chat_images(chat_name: str, limit: int = 20) -> str: +def get_chat_images(chat_name: str, limit: int = 20, offset: int = 0, start_time: str = "", end_time: str = "") -> str: """列出某个聊天中的图片消息。 返回图片的时间、local_id、MD5、文件大小等信息。 @@ -2839,7 +2839,16 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str: Args: chat_name: 聊天对象的名字、备注名或wxid limit: 返回数量,默认20 + offset: 分页偏移量,默认0 + start_time: 起始时间,支持 YYYY-MM-DD / YYYY-MM-DD HH:MM / YYYY-MM-DD HH:MM:SS + end_time: 结束时间,支持 YYYY-MM-DD / YYYY-MM-DD HH:MM / YYYY-MM-DD HH:MM:SS """ + try: + _validate_pagination(limit, offset) + start_ts, end_ts = _parse_time_range(start_time, end_time) + except ValueError as e: + return f"错误: {e}" + username = resolve_username(chat_name) if not username: return f"找不到聊天对象: {chat_name}" @@ -2854,12 +2863,15 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str: if not shards: return f"找不到 {display_name} 的消息记录" - # 每个 shard 取 limit 张, 合并后按 create_time DESC 全局排序, 取最新 limit 张。 - # 单 shard 至少够本次返回, 避免一个 shard 凑不出 limit 时其他 shard 没机会贡献。 + # 每个 shard 取 limit+offset 张候选, 合并后按 create_time DESC 全局排序, 切片 + # [offset : offset+limit] 出本页。单 shard 至少凑得起本页, 避免某 shard 缺数据 + # 时本页变短。 + candidate_limit = limit + offset all_images = [] for shard in shards: shard_images = _image_resolver.list_chat_images( - shard['db_path'], shard['table_name'], username, limit + shard['db_path'], shard['table_name'], username, + limit=candidate_limit, start_ts=start_ts, end_ts=end_ts, ) all_images.extend(shard_images) @@ -2867,10 +2879,10 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str: return f"{display_name} 无图片消息" all_images.sort(key=lambda img: img['create_time'], reverse=True) - images = all_images[:limit] + paged = all_images[offset:offset + limit] lines = [] - for img in images: + for img in paged: time_str = datetime.fromtimestamp(img['create_time']).strftime('%Y-%m-%d %H:%M') line = f"[{time_str}] local_id={img['local_id']}" if img.get('md5'): @@ -2882,7 +2894,10 @@ def get_chat_images(chat_name: str, limit: int = 20) -> str: line += " (无资源信息)" lines.append(line) - return f"{display_name} 的 {len(lines)} 张图片:\n\n" + "\n".join(lines) + _pagination_hint(len(lines), limit, 0) + header = f"{display_name} 的 {len(lines)} 张图片(offset={offset}, limit={limit})" + if start_time or end_time: + header += f"\n时间范围: {start_time or '最早'} ~ {end_time or '最新'}" + return header + ":\n\n" + "\n".join(lines) + _pagination_hint(len(lines), limit, offset) # ============ 语音解密 ============ @@ -2952,7 +2967,7 @@ def _silk_to_wav(voice_data, create_time, username, local_id): @mcp.tool() -def get_voice_messages(chat_name: str, limit: int = 20) -> str: +def get_voice_messages(chat_name: str, limit: int = 20, offset: int = 0, start_time: str = "", end_time: str = "") -> str: """列出某个聊天中的语音消息。 返回语音的时间、local_id 和大小,可配合 decode_voice 工具解码。 @@ -2960,7 +2975,16 @@ def get_voice_messages(chat_name: str, limit: int = 20) -> str: Args: chat_name: 聊天对象的名字、备注名或wxid limit: 返回数量,默认20 + offset: 分页偏移量,默认0 + start_time: 起始时间,支持 YYYY-MM-DD / YYYY-MM-DD HH:MM / YYYY-MM-DD HH:MM:SS + end_time: 结束时间,支持 YYYY-MM-DD / YYYY-MM-DD HH:MM / YYYY-MM-DD HH:MM:SS """ + try: + _validate_pagination(limit, offset) + start_ts, end_ts = _parse_time_range(start_time, end_time) + except ValueError as e: + return f"错误: {e}" + username = resolve_username(chat_name) if not username: return f"找不到聊天对象: {chat_name}" @@ -2971,31 +2995,48 @@ def get_voice_messages(chat_name: str, limit: int = 20) -> str: if not MEDIA_DB_KEYS: return "找不到 media DB" - # 从每个分片各取最多 limit 条后合并再截断:分片若有时间重叠也不会漏最新消息 + # 每分片各取 limit+offset 条候选, 合并后全局排序切片 [offset:offset+limit] 出本页。 + candidate_limit = limit + offset + clauses = ['chat_name_id = ?'] + if start_ts is not None: + clauses.append('create_time >= ?') + if end_ts is not None: + clauses.append('create_time <= ?') + where_sql = ' AND '.join(clauses) + rows = [] for media_db in _iter_media_db_paths(): with closing(sqlite3.connect(media_db)) as conn: chat_name_id = _get_chat_name_id(conn, username) if chat_name_id is None: continue + params = [chat_name_id] + if start_ts is not None: + params.append(start_ts) + if end_ts is not None: + params.append(end_ts) + params.append(candidate_limit) rows.extend(conn.execute( - "SELECT local_id, create_time, length(voice_data) FROM VoiceInfo " - "WHERE chat_name_id = ? ORDER BY create_time DESC LIMIT ?", - (chat_name_id, limit), + f"SELECT local_id, create_time, length(voice_data) FROM VoiceInfo " + f"WHERE {where_sql} ORDER BY create_time DESC LIMIT ?", + params, ).fetchall()) if not rows: return f"{display_name} 无语音消息" rows.sort(key=lambda r: r[1], reverse=True) - rows = rows[:limit] + paged = rows[offset:offset + limit] lines = [] - for local_id, create_time, size in rows: + for local_id, create_time, size in paged: time_str = datetime.fromtimestamp(create_time).strftime('%Y-%m-%d %H:%M') lines.append(f"[{time_str}] local_id={local_id} {size/1024:.0f}KB") - return f"{display_name} 的 {len(lines)} 条语音消息:\n\n" + "\n".join(lines) + _pagination_hint(len(lines), limit, 0) + header = f"{display_name} 的 {len(lines)} 条语音消息(offset={offset}, limit={limit})" + if start_time or end_time: + header += f"\n时间范围: {start_time or '最早'} ~ {end_time or '最新'}" + return header + ":\n\n" + "\n".join(lines) + _pagination_hint(len(lines), limit, offset) @mcp.tool() diff --git a/tests/test_chat_images_query_align.py b/tests/test_chat_images_query_align.py new file mode 100644 index 0000000..c6cca2f --- /dev/null +++ b/tests/test_chat_images_query_align.py @@ -0,0 +1,97 @@ +"""测试 get_chat_images 新增的 offset / start_time / end_time 参数。""" +import os +import sys +from unittest.mock import patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import mcp_server + + +def _img(local_id, create_time, md5=None, size=None): + info = {'local_id': local_id, 'create_time': create_time, 'md5': md5} + if size is not None: + info['size'] = size + return info + + +def _run_with(shard_images_map, **kwargs): + """Helper: stub collaborators and call get_chat_images with new kwargs.""" + shards = [{'db_path': k, 'table_name': 'Msg_x'} for k in shard_images_map] + captured = {'calls': []} + + def fake_list(db_path, table_name, username, limit=20, start_ts=None, end_ts=None): + captured['calls'].append({ + 'db_path': db_path, 'limit': limit, 'start_ts': start_ts, 'end_ts': end_ts, + }) + return shard_images_map.get(db_path, []) + + with patch.object(mcp_server, 'resolve_username', return_value='wxid_demo'), \ + patch.object(mcp_server, 'get_contact_names', return_value={'wxid_demo': 'Demo'}), \ + patch.object(mcp_server, '_find_msg_tables_for_user', return_value=shards), \ + patch.object(mcp_server._image_resolver, 'list_chat_images', side_effect=fake_list): + return mcp_server.get_chat_images('Demo', **kwargs), captured + + +def test_invalid_offset_returns_error(): + out, _ = _run_with({}, offset=-1) + assert '错误' in out + + +def test_invalid_time_range_returns_error(): + """start_time 晚于 end_time 应报错。""" + out, _ = _run_with({}, start_time='2026-05-10', end_time='2026-05-01') + assert '错误' in out + + +def test_candidate_limit_includes_offset(): + """每 shard 拉 limit+offset 张候选, 保证全局分页能切到正确的页。""" + _, captured = _run_with({'/a': [_img(1, 1000)]}, limit=5, offset=10) + assert captured['calls'][0]['limit'] == 15 + + +def test_start_end_ts_forwarded_to_shard_query(): + """start_time / end_time 解析为 unix 秒后透传给 shard 查询。""" + _, captured = _run_with( + {'/a': []}, + start_time='2026-05-01', + end_time='2026-05-31', + ) + call = captured['calls'][0] + assert call['start_ts'] is not None + assert call['end_ts'] is not None + assert call['start_ts'] < call['end_ts'] + + +def test_offset_slices_paged_window(): + """offset=2, limit=2 取全局排序后第 3-4 张图片。""" + shard_a = [_img(1, 1100), _img(2, 1000)] + shard_b = [_img(3, 1300), _img(4, 1200)] + out, _ = _run_with({'/a': shard_a, '/b': shard_b}, limit=2, offset=2) + # 全局排序后顺序: 1300, 1200, 1100, 1000 → 第 3-4 是 1100, 1000 → local_id 1, 2 + assert 'local_id=1' in out + assert 'local_id=2' in out + assert 'local_id=3' not in out + assert 'local_id=4' not in out + + +def test_header_shows_time_range_when_given(): + shard_a = [_img(1, 1000, md5='abc')] + out, _ = _run_with({'/a': shard_a}, start_time='2026-05-01') + assert '时间范围' in out + assert '2026-05-01' in out + + +def test_header_shows_offset_limit(): + shard_a = [_img(1, 1000, md5='abc')] + out, _ = _run_with({'/a': shard_a}, limit=10, offset=20) + assert 'offset=20' in out + assert 'limit=10' in out + + +def test_default_behavior_unchanged(): + """不传新参数时行为与旧接口一致 — offset=0 切片就是 [:limit]。""" + shard_a = [_img(1, 1100, md5='a1'), _img(2, 1000, md5='a2')] + out, _ = _run_with({'/a': shard_a}) + assert 'local_id=1' in out + assert 'local_id=2' in out diff --git a/tests/test_get_chat_images_multishard.py b/tests/test_get_chat_images_multishard.py index 2a2f589..75afa09 100644 --- a/tests/test_get_chat_images_multishard.py +++ b/tests/test_get_chat_images_multishard.py @@ -33,7 +33,7 @@ class GetChatImagesMultiShardTests(unittest.TestCase): def _run(self, shards, shard_images_map, limit=20): """Helper: stub the two collaborators and call the tool.""" - def fake_list(db_path, table_name, username, lim): + def fake_list(db_path, table_name, username, limit=20, start_ts=None, end_ts=None): return shard_images_map.get(db_path, []) with patch.object(mcp_server, "_find_msg_tables_for_user",