diff --git a/api/apps/restful_apis/memory_api.py b/api/apps/restful_apis/memory_api.py index b455fca1f8..763befeb3d 100644 --- a/api/apps/restful_apis/memory_api.py +++ b/api/apps/restful_apis/memory_api.py @@ -260,7 +260,10 @@ async def search_message(): query = args.get("query") similarity_threshold = float(args.get("similarity_threshold", 0.2)) keywords_similarity_weight = float(args.get("keywords_similarity_weight", 0.7)) - top_n = int(args.get("top_n", 5)) + try: + top_n = validate_rest_api_page_size(int(args.get("top_n", 5))) + except ValueError as exc: + return get_error_argument_result(str(exc)) agent_id = args.get("agent_id", "") session_id = args.get("session_id", "") user_id = args.get("user_id", "") @@ -280,7 +283,10 @@ async def get_messages(): memory_ids = memory_ids[0].split(",") agent_id = args.get("agent_id", "") session_id = args.get("session_id", "") - limit = int(args.get("limit", 10)) + try: + limit = validate_rest_api_page_size(int(args.get("limit", 10))) + except ValueError as exc: + return get_error_argument_result(str(exc)) if not memory_ids: return get_error_argument_result("memory_ids is required.") try: diff --git a/test/testcases/test_web_api/test_message_app/test_message_routes_unit.py b/test/testcases/test_web_api/test_message_app/test_message_routes_unit.py index f4641ed469..820c9782f9 100644 --- a/test/testcases/test_web_api/test_message_app/test_message_routes_unit.py +++ b/test/testcases/test_web_api/test_message_app/test_message_routes_unit.py @@ -149,3 +149,36 @@ def test_get_messages_csv_and_missing_memory_ids(monkeypatch): res = _run(inspect.unwrap(module.get_messages)()) assert res["code"] == module.RetCode.SUCCESS, res assert isinstance(res["data"], list), res + + +@pytest.mark.p2 +def test_get_messages_rejects_limit_above_rest_api_max(monkeypatch): + module = _load_memory_routes_module(monkeypatch) + monkeypatch.setattr( + module, + "request", + SimpleNamespace(args=_DummyArgs({"memory_id": "m1", "limit": "1000"})), + ) + res = _run(inspect.unwrap(module.get_messages)()) + assert res["code"] == module.RetCode.ARGUMENT_ERROR, res + assert "page_size must be less than or equal to 100" in res["message"], res + + +@pytest.mark.p2 +def test_search_message_rejects_top_n_above_rest_api_max(monkeypatch): + module = _load_memory_routes_module(monkeypatch) + monkeypatch.setattr( + module, + "request", + SimpleNamespace( + args=_DummyArgs({"memory_id": "m1", "query": "hello", "top_n": "500"}) + ), + ) + + async def _search_message(_filter_dict, _params): + return [] + + monkeypatch.setattr(module.memory_api_service, "search_message", _search_message) + res = _run(inspect.unwrap(module.search_message)()) + assert res["code"] == module.RetCode.ARGUMENT_ERROR, res + assert "page_size must be less than or equal to 100" in res["message"], res