From fc3938e954a5d173cadb85da4fe3895d83ae1aaf Mon Sep 17 00:00:00 2001 From: nankingjing <1079826437@qq.com> Date: Sun, 12 Jul 2026 08:51:21 +0800 Subject: [PATCH 1/2] fix(retrieve): only count preference_full_count on actual full render The preference_full_count counter was incremented before the budget check, so the first PREFERENCE_FULL_LIMIT preference memories that *failed* the budget check (e.g. oversized content) wasted the cap. Subsequent preferences that would have fit were denied full rendering even though budget was still available. Move the increment inside the successful-full branch so it only counts when a preference actually renders as 'full'. Add a regression test verifying that preferences after the limit can still render full when the preceding ones overflowed the budget. --- openviking/retrieve/type_quota_recall.py | 3 +- .../test_type_quota_preference_counter.py | 175 ++++++++++++++++++ 2 files changed, 177 insertions(+), 1 deletion(-) create mode 100644 tests/server/test_type_quota_preference_counter.py diff --git a/openviking/retrieve/type_quota_recall.py b/openviking/retrieve/type_quota_recall.py index a657f27c2..402fd3e45 100644 --- a/openviking/retrieve/type_quota_recall.py +++ b/openviking/retrieve/type_quota_recall.py @@ -405,7 +405,6 @@ async def search_type_quota_recall( can_try_full = memory_type in budgets if memory_type == "preferences": can_try_full = preference_full_count < PREFERENCE_FULL_LIMIT - preference_full_count += 1 if ( can_try_full and used_by_type.get(memory_type, 0) + full_chars @@ -417,6 +416,8 @@ async def search_type_quota_recall( fragment = full used_by_type[memory_type] = used_by_type.get(memory_type, 0) + full_chars total_chars += full_chars + if memory_type == "preferences": + preference_full_count += 1 elif memory_type == "events": summary = _extract_event_summary(content, fallback=abstract) if summary: diff --git a/tests/server/test_type_quota_preference_counter.py b/tests/server/test_type_quota_preference_counter.py new file mode 100644 index 000000000..058effe71 --- /dev/null +++ b/tests/server/test_type_quota_preference_counter.py @@ -0,0 +1,175 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +"""Regression tests for type_quota_recall preference_full_count behavior. + +Bug #8: ``preference_full_count`` was incremented before the budget check, +so the first ``PREFERENCE_FULL_LIMIT`` preferences that *failed* the budget +check wasted the cap. Subsequent preferences that *would* have fit were +denied full rendering. + +These tests are standalone unit tests (no server fixtures needed) so they +run without a full OpenViking deployment or config file. +""" + +from types import SimpleNamespace + +import pytest + +from openviking.retrieve.type_quota_recall import ( + PREFERENCE_FULL_LIMIT, + search_type_quota_recall, +) +from openviking.server.identity import RequestContext, Role +from openviking_cli.retrieve import ContextType, MatchedContext +from openviking_cli.session.user_id import UserIdentifier + + +class _FakeFindResult: + def __init__(self, memories): + self.memories = memories + + +def _make_mock_service(find_fn, read_fn): + """Build a minimal mock service with search.find and fs.read.""" + search = SimpleNamespace(find=find_fn) + fs = SimpleNamespace(read=read_fn) + return SimpleNamespace(search=search, fs=fs) + + +@pytest.mark.asyncio +async def test_preference_full_count_only_incremented_on_full_success(): + """``preference_full_count`` must only increment when a preference actually + renders as *full*, not before the budget check. + + Scenario: ``PREFERENCE_FULL_LIMIT + 2`` preferences returned. The first + ``PREFERENCE_FULL_LIMIT`` entries have content so large that their full + fragments exceed ``max_chars``. The remaining two entries have small + content. After the fix the last two entries should still render as + ``"full"`` — the counter was not wasted on the first three (failing) + entries. + """ + pref_count = PREFERENCE_FULL_LIMIT + 2 + mem_root = "viking://user/test_user/memories" + + async def fake_find(**kwargs): + target = kwargs["target_uri"] + if target.endswith("/preferences"): + return _FakeFindResult( + [ + MatchedContext( + uri=f"{mem_root}/preferences/pref_{i}.md", + context_type=ContextType.MEMORY, + level=2, + score=0.9 - i * 0.01, + abstract=f"preference {i}", + category="preferences", + ) + for i in range(pref_count) + ] + ) + return _FakeFindResult([]) + + read_contents: dict[str, str] = {} + for i in range(pref_count): + if i < PREFERENCE_FULL_LIMIT: + # Very large content — its full fragment will exceed the budget. + read_contents[f"{mem_root}/preferences/pref_{i}.md"] = "x" * 5000 + else: + read_contents[f"{mem_root}/preferences/pref_{i}.md"] = "short pref" + + async def fake_read(uri, **kwargs): + return read_contents.get(uri, "") + + service = _make_mock_service(fake_find, fake_read) + + ctx = RequestContext( + user=UserIdentifier.the_default_user("test_user"), + role=Role.ROOT, + ) + + result = await search_type_quota_recall( + service=service, + ctx=ctx, + query="preferences test", + quotas={"events": 0, "entities": 0, "preferences": pref_count, "experiences": 0}, + max_chars=500, + min_score=0.1, + render=True, + ) + + # The first PREFERENCE_FULL_LIMIT entries should NOT be "full" — content + # was too large. + for i in range(PREFERENCE_FULL_LIMIT): + assert result.entries[i].mode != "full", ( + f"Entry {i} with large content unexpectedly rendered as full" + ) + + # Entries PREFERENCE_FULL_LIMIT and beyond that DO fit the budget should + # still render as "full" — the counter wasn't wasted on the failures. + for i in range(PREFERENCE_FULL_LIMIT, pref_count): + assert result.entries[i].mode == "full", ( + f"Entry {i} should be full but was {result.entries[i].mode!r}. " + f"Modes: {[e.mode for e in result.entries]}" + ) + + +@pytest.mark.asyncio +async def test_preference_full_limit_still_capped_when_full_fits(): + """After ``PREFERENCE_FULL_LIMIT`` preferences successfully render as full, + subsequent preferences must be capped (no longer try full) even if budget + remains. + """ + pref_count = PREFERENCE_FULL_LIMIT + 3 + mem_root = "viking://user/test_user/memories" + + async def fake_find(**kwargs): + target = kwargs["target_uri"] + if target.endswith("/preferences"): + return _FakeFindResult( + [ + MatchedContext( + uri=f"{mem_root}/preferences/pref_{i}.md", + context_type=ContextType.MEMORY, + level=2, + score=0.9 - i * 0.01, + abstract=f"preference {i}", + category="preferences", + ) + for i in range(pref_count) + ] + ) + return _FakeFindResult([]) + + async def fake_read(uri, **kwargs): + return "short content" + + service = _make_mock_service(fake_find, fake_read) + + ctx = RequestContext( + user=UserIdentifier.the_default_user("test_user"), + role=Role.ROOT, + ) + + result = await search_type_quota_recall( + service=service, + ctx=ctx, + query="preferences test", + quotas={"events": 0, "entities": 0, "preferences": pref_count, "experiences": 0}, + max_chars=10000, + min_score=0.1, + render=True, + ) + + full_modes = [e.mode for e in result.entries if e.mode == "full"] + assert len(full_modes) == PREFERENCE_FULL_LIMIT, ( + f"Expected exactly {PREFERENCE_FULL_LIMIT} full preferences, " + f"got {len(full_modes)}. Modes: {[e.mode for e in result.entries]}" + ) + + # After the limit, remaining entries should NOT be "full" + for i in range(PREFERENCE_FULL_LIMIT, len(result.entries)): + assert result.entries[i].mode != "full", ( + f"Entry {i} should not be full (limit reached). " + f"Got mode={result.entries[i].mode!r}" + ) From 49361b07a27ca9093b0b1e64f21e90646813dad1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=BB=84=E4=BA=91=E9=BE=99?= <76432572+nankingjing@users.noreply.github.com> Date: Tue, 14 Jul 2026 09:34:19 +0800 Subject: [PATCH 2/2] fix(tests): adapt preference_full_count regression to upstream dedupe and parallelization --- .../test_type_quota_preference_counter.py | 362 +++++++++--------- 1 file changed, 187 insertions(+), 175 deletions(-) diff --git a/tests/server/test_type_quota_preference_counter.py b/tests/server/test_type_quota_preference_counter.py index 058effe71..20c9b2ef8 100644 --- a/tests/server/test_type_quota_preference_counter.py +++ b/tests/server/test_type_quota_preference_counter.py @@ -1,175 +1,187 @@ -# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. -# SPDX-License-Identifier: AGPL-3.0 - -"""Regression tests for type_quota_recall preference_full_count behavior. - -Bug #8: ``preference_full_count`` was incremented before the budget check, -so the first ``PREFERENCE_FULL_LIMIT`` preferences that *failed* the budget -check wasted the cap. Subsequent preferences that *would* have fit were -denied full rendering. - -These tests are standalone unit tests (no server fixtures needed) so they -run without a full OpenViking deployment or config file. -""" - -from types import SimpleNamespace - -import pytest - -from openviking.retrieve.type_quota_recall import ( - PREFERENCE_FULL_LIMIT, - search_type_quota_recall, -) -from openviking.server.identity import RequestContext, Role -from openviking_cli.retrieve import ContextType, MatchedContext -from openviking_cli.session.user_id import UserIdentifier - - -class _FakeFindResult: - def __init__(self, memories): - self.memories = memories - - -def _make_mock_service(find_fn, read_fn): - """Build a minimal mock service with search.find and fs.read.""" - search = SimpleNamespace(find=find_fn) - fs = SimpleNamespace(read=read_fn) - return SimpleNamespace(search=search, fs=fs) - - -@pytest.mark.asyncio -async def test_preference_full_count_only_incremented_on_full_success(): - """``preference_full_count`` must only increment when a preference actually - renders as *full*, not before the budget check. - - Scenario: ``PREFERENCE_FULL_LIMIT + 2`` preferences returned. The first - ``PREFERENCE_FULL_LIMIT`` entries have content so large that their full - fragments exceed ``max_chars``. The remaining two entries have small - content. After the fix the last two entries should still render as - ``"full"`` — the counter was not wasted on the first three (failing) - entries. - """ - pref_count = PREFERENCE_FULL_LIMIT + 2 - mem_root = "viking://user/test_user/memories" - - async def fake_find(**kwargs): - target = kwargs["target_uri"] - if target.endswith("/preferences"): - return _FakeFindResult( - [ - MatchedContext( - uri=f"{mem_root}/preferences/pref_{i}.md", - context_type=ContextType.MEMORY, - level=2, - score=0.9 - i * 0.01, - abstract=f"preference {i}", - category="preferences", - ) - for i in range(pref_count) - ] - ) - return _FakeFindResult([]) - - read_contents: dict[str, str] = {} - for i in range(pref_count): - if i < PREFERENCE_FULL_LIMIT: - # Very large content — its full fragment will exceed the budget. - read_contents[f"{mem_root}/preferences/pref_{i}.md"] = "x" * 5000 - else: - read_contents[f"{mem_root}/preferences/pref_{i}.md"] = "short pref" - - async def fake_read(uri, **kwargs): - return read_contents.get(uri, "") - - service = _make_mock_service(fake_find, fake_read) - - ctx = RequestContext( - user=UserIdentifier.the_default_user("test_user"), - role=Role.ROOT, - ) - - result = await search_type_quota_recall( - service=service, - ctx=ctx, - query="preferences test", - quotas={"events": 0, "entities": 0, "preferences": pref_count, "experiences": 0}, - max_chars=500, - min_score=0.1, - render=True, - ) - - # The first PREFERENCE_FULL_LIMIT entries should NOT be "full" — content - # was too large. - for i in range(PREFERENCE_FULL_LIMIT): - assert result.entries[i].mode != "full", ( - f"Entry {i} with large content unexpectedly rendered as full" - ) - - # Entries PREFERENCE_FULL_LIMIT and beyond that DO fit the budget should - # still render as "full" — the counter wasn't wasted on the failures. - for i in range(PREFERENCE_FULL_LIMIT, pref_count): - assert result.entries[i].mode == "full", ( - f"Entry {i} should be full but was {result.entries[i].mode!r}. " - f"Modes: {[e.mode for e in result.entries]}" - ) - - -@pytest.mark.asyncio -async def test_preference_full_limit_still_capped_when_full_fits(): - """After ``PREFERENCE_FULL_LIMIT`` preferences successfully render as full, - subsequent preferences must be capped (no longer try full) even if budget - remains. - """ - pref_count = PREFERENCE_FULL_LIMIT + 3 - mem_root = "viking://user/test_user/memories" - - async def fake_find(**kwargs): - target = kwargs["target_uri"] - if target.endswith("/preferences"): - return _FakeFindResult( - [ - MatchedContext( - uri=f"{mem_root}/preferences/pref_{i}.md", - context_type=ContextType.MEMORY, - level=2, - score=0.9 - i * 0.01, - abstract=f"preference {i}", - category="preferences", - ) - for i in range(pref_count) - ] - ) - return _FakeFindResult([]) - - async def fake_read(uri, **kwargs): - return "short content" - - service = _make_mock_service(fake_find, fake_read) - - ctx = RequestContext( - user=UserIdentifier.the_default_user("test_user"), - role=Role.ROOT, - ) - - result = await search_type_quota_recall( - service=service, - ctx=ctx, - query="preferences test", - quotas={"events": 0, "entities": 0, "preferences": pref_count, "experiences": 0}, - max_chars=10000, - min_score=0.1, - render=True, - ) - - full_modes = [e.mode for e in result.entries if e.mode == "full"] - assert len(full_modes) == PREFERENCE_FULL_LIMIT, ( - f"Expected exactly {PREFERENCE_FULL_LIMIT} full preferences, " - f"got {len(full_modes)}. Modes: {[e.mode for e in result.entries]}" - ) - - # After the limit, remaining entries should NOT be "full" - for i in range(PREFERENCE_FULL_LIMIT, len(result.entries)): - assert result.entries[i].mode != "full", ( - f"Entry {i} should not be full (limit reached). " - f"Got mode={result.entries[i].mode!r}" - ) +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +"""Regression tests for type_quota_recall preference_full_count behavior. + +Bug #8: ``preference_full_count`` was incremented before the budget check, +so the first ``PREFERENCE_FULL_LIMIT`` preferences that *failed* the budget +check wasted the cap. Subsequent preferences that *would* have fit were +denied full rendering. + +These tests are standalone unit tests (no server fixtures needed) so they +run without a full OpenViking deployment or config file. +""" + +from types import SimpleNamespace + +import pytest + +from openviking.retrieve.type_quota_recall import ( + PREFERENCE_FULL_LIMIT, + search_type_quota_recall, +) +from openviking.server.identity import RequestContext, Role +from openviking_cli.retrieve import ContextType, MatchedContext +from openviking_cli.session.user_id import UserIdentifier + + +class _FakeFindResult: + def __init__(self, memories): + self.memories = memories + + +def _make_mock_service(find_fn, read_fn): + """Build a minimal mock service with search.find and fs.read.""" + search = SimpleNamespace(find=find_fn) + fs = SimpleNamespace(read=read_fn) + return SimpleNamespace(search=search, fs=fs) + + +@pytest.mark.asyncio +async def test_preference_full_count_only_incremented_on_full_success(): + """``preference_full_count`` must only increment when a preference actually + renders as *full*, not before the budget check. + + Scenario: ``PREFERENCE_FULL_LIMIT + 2`` preferences returned. The first + ``PREFERENCE_FULL_LIMIT`` entries have content so large that their full + fragments exceed ``max_chars``. The remaining two entries have small + content. After the fix the last two entries should still render as + ``"full"`` — the counter was not wasted on the first three (failing) + entries. + """ + pref_count = PREFERENCE_FULL_LIMIT + 2 + mem_root = "viking://user/test_user/memories" + + async def fake_find(**kwargs): + target = kwargs["target_uri"] + if target.endswith("/preferences"): + return _FakeFindResult( + [ + MatchedContext( + uri=f"{mem_root}/preferences/pref_{i}.md", + context_type=ContextType.MEMORY, + level=2, + score=0.9 - i * 0.01, + abstract=f"preference {i}", + category="preferences", + ) + for i in range(pref_count) + ] + ) + return _FakeFindResult([]) + + read_contents: dict[str, str] = {} + for i in range(pref_count): + if i < PREFERENCE_FULL_LIMIT: + # Very large content — its full fragment will exceed the budget. + # Make it unique per memory so the new content-hash dedupe does not + # collapse the failing entries together. + read_contents[f"{mem_root}/preferences/pref_{i}.md"] = f"large content block {i} " + "x" * 5000 + else: + read_contents[f"{mem_root}/preferences/pref_{i}.md"] = f"short pref {i}" + + async def fake_read(uri, **kwargs): + return read_contents.get(uri, "") + + service = _make_mock_service(fake_find, fake_read) + + ctx = RequestContext( + user=UserIdentifier.the_default_user("test_user"), + role=Role.ROOT, + ) + + result = await search_type_quota_recall( + service=service, + ctx=ctx, + query="preferences test", + quotas={"events": 0, "entities": 0, "preferences": pref_count, "experiences": 0}, + max_chars=2000, + min_score=0.1, + render=True, + ) + + by_uri = {e.uri: e for e in result.entries} + large_uris = {f"{mem_root}/preferences/pref_{i}.md" for i in range(PREFERENCE_FULL_LIMIT)} + small_uris = {f"{mem_root}/preferences/pref_{i}.md" for i in range(PREFERENCE_FULL_LIMIT, pref_count)} + + # Large-content entries should not have rendered as full; they may still + # appear as URI-only entries (the upstream budget fallback). + for uri in large_uris: + entry = by_uri.get(uri) + assert entry is not None, f"Missing large-content entry {uri}. Got {list(by_uri)}" + assert entry.mode != "full", ( + f"Large-content entry {uri} unexpectedly rendered as full" + ) + + # Small-content entries that fit must still render as full: the counter + # was not consumed by the failing large-content entries. + for uri in small_uris: + entry = by_uri.get(uri) + assert entry is not None, f"Missing small-content entry {uri}. Got {list(by_uri)}" + assert entry.mode == "full", ( + f"Entry {uri} should be full but was {entry.mode!r}. " + f"Modes: {[e.mode for e in result.entries]}" + ) + + +@pytest.mark.asyncio +async def test_preference_full_limit_still_capped_when_full_fits(): + """After ``PREFERENCE_FULL_LIMIT`` preferences successfully render as full, + subsequent preferences must be capped (no longer try full) even if budget + remains. + """ + pref_count = PREFERENCE_FULL_LIMIT + 3 + mem_root = "viking://user/test_user/memories" + + async def fake_find(**kwargs): + target = kwargs["target_uri"] + if target.endswith("/preferences"): + return _FakeFindResult( + [ + MatchedContext( + uri=f"{mem_root}/preferences/pref_{i}.md", + context_type=ContextType.MEMORY, + level=2, + score=0.9 - i * 0.01, + abstract=f"preference {i}", + category="preferences", + ) + for i in range(pref_count) + ] + ) + return _FakeFindResult([]) + + async def fake_read(uri, **kwargs): + # Unique per-memory content so the upstream content-hash deduper + # does not collapse distinct preferences. + return f"short content {uri}" + + service = _make_mock_service(fake_find, fake_read) + + ctx = RequestContext( + user=UserIdentifier.the_default_user("test_user"), + role=Role.ROOT, + ) + + result = await search_type_quota_recall( + service=service, + ctx=ctx, + query="preferences test", + quotas={"events": 0, "entities": 0, "preferences": pref_count, "experiences": 0}, + max_chars=10000, + min_score=0.1, + render=True, + ) + + full_modes = [e.mode for e in result.entries if e.mode == "full"] + assert len(full_modes) == PREFERENCE_FULL_LIMIT, ( + f"Expected exactly {PREFERENCE_FULL_LIMIT} full preferences, " + f"got {len(full_modes)}. Modes: {[e.mode for e in result.entries]}" + ) + + # After the limit, remaining entries should NOT be "full" + for i in range(PREFERENCE_FULL_LIMIT, len(result.entries)): + assert result.entries[i].mode != "full", ( + f"Entry {i} should not be full (limit reached). " + f"Got mode={result.entries[i].mode!r}" + )