From b5a59145593c1b4981a5f01096be7ef90b940862 Mon Sep 17 00:00:00 2001 From: Eurekaxun Date: Tue, 4 Aug 2026 14:44:09 +0800 Subject: [PATCH] fix(memory): reject pseudo tool outputs during extraction --- .../session_extract_context_provider.py | 36 +++++++- openviking/session/memory/tools.py | 10 +- .../session/memory/utils/json_parser.py | 2 + tests/session/memory/test_json_stability.py | 25 +++++ tests/session/memory/test_memory_react.py | 92 ++++++++++++++++++- .../memory/test_memory_react_system_prompt.py | 46 ++++++++-- tests/session/memory/test_memory_tools.py | 20 ++++ .../test_patch_merge_context_provider.py | 4 +- 8 files changed, 216 insertions(+), 19 deletions(-) diff --git a/openviking/session/memory/session_extract_context_provider.py b/openviking/session/memory/session_extract_context_provider.py index a7990234de..a9a6dc4772 100644 --- a/openviking/session/memory/session_extract_context_provider.py +++ b/openviking/session/memory/session_extract_context_provider.py @@ -207,17 +207,43 @@ def instruction(self) -> str: if contains_resource_uri else "" ) + if self._eager_prefetch: + resource_deletion_prefetch_rule = ( + "\n- For URIs listed under the system-generated `## Resource Deletion` block's " + "`Affected memory URIs`, only edit files whose complete content is already in the " + "pre-fetched context" + if contains_resource_uri + else "" + ) + context_workflow = ( + "2. Use only the pre-fetched context; no tools are available\n" + "3. Output ONLY a JSON object (no extra text before or after)" + ) + tool_rules = ( + "- No tools are available. Do not output read, search, write, or other tool requests\n" + "- Only edit an existing memory file when its complete content is already included " + f"in the pre-fetched context{resource_deletion_prefetch_rule}" + ) + else: + context_workflow = ( + "2. If you need the complete content of a listed memory URI, use the read tool\n" + "3. When you have enough information, output ONLY a JSON object " + "(no extra text before or after)" + ) + tool_rules = ( + "- ONLY the read tool is available - search and write are not available\n" + "- Before editing ANY existing memory file, you MUST first read its complete content\n" + "- ONLY read URIs that are explicitly listed in pre-fetched search results, " + f"returned by previous tool calls{resource_deletion_read_source}" + ) goal = f"""You are a memory extraction agent. Your task is to analyze conversations and update memories. ## Workflow 1. Analyze the conversation and pre-fetched context -2. If you need more information, use the available tools (read/search) -3. When you have enough information, output ONLY a JSON object (no extra text before or after) +{context_workflow} ## Critical -- ONLY read and search tools are available - DO NOT use write tool -- Before editing ANY existing memory file, you MUST first read its complete content -- ONLY read URIs that are explicitly listed in ls/search tool results, returned by previous tool calls{resource_deletion_read_source} +{tool_rules} ## Target Output Language All memory content MUST be written in {output_language}. diff --git a/openviking/session/memory/tools.py b/openviking/session/memory/tools.py index 76dee7e9ec..86c904161c 100644 --- a/openviking/session/memory/tools.py +++ b/openviking/session/memory/tools.py @@ -75,12 +75,18 @@ def add_tool_call_pair_to_messages( params: Dict[str, Any], result: Any, ) -> None: - """Add a tool call pair with optimized format to save tokens.""" + """Add a compact tool result message without imitating a callable request.""" messages.append( { "role": "user", "content": json.dumps( - {"tool_call_name": tool_name, "args": params, "result": result}, ensure_ascii=False + { + "message_type": "tool_result", + "tool_name": tool_name, + "args": params, + "result": result, + }, + ensure_ascii=False, ), } ) diff --git a/openviking/session/memory/utils/json_parser.py b/openviking/session/memory/utils/json_parser.py index 380a4f0933..10340b1213 100644 --- a/openviking/session/memory/utils/json_parser.py +++ b/openviking/session/memory/utils/json_parser.py @@ -430,6 +430,8 @@ def parse_json_with_stability( # Filter to only expected fields if provided if expected_fields: + if parsed_data and not any(key in expected_fields for key in parsed_data): + return None, "No recognized fields in non-empty JSON object" filtered_data = {} for k, v in parsed_data.items(): if k in expected_fields: diff --git a/tests/session/memory/test_json_stability.py b/tests/session/memory/test_json_stability.py index c76072da7b..1d76a095c2 100644 --- a/tests/session/memory/test_json_stability.py +++ b/tests/session/memory/test_json_stability.py @@ -206,6 +206,31 @@ def test_filters_extra_fields(self): assert data.reasonning == "test" assert data.count == 42 + def test_rejects_nonempty_object_without_any_expected_field(self): + """Pseudo tool calls must not become successful empty operations.""" + content = '{"tool_call_name":"search","args":{"query":"Melanie pottery"}}' + + data, error = parse_json_with_stability( + content, + model_class=self.TestModel, + expected_fields=["reasonning", "count", "tags"], + ) + + assert data is None + assert error == "No recognized fields in non-empty JSON object" + + def test_keeps_explicit_empty_object_for_operations_compatibility(self): + """An explicit empty object remains a valid no-operations response.""" + data, error = parse_json_with_stability( + "{}", + model_class=self.TestModel, + expected_fields=["reasonning", "count", "tags"], + ) + + assert error is None + assert data.reasonning == "" + assert data.tags == [] + def test_returns_raw_dict_when_no_model_class(self): """Test returns dict when no model_class is provided.""" content = '{"reasonning": "test"}' diff --git a/tests/session/memory/test_memory_react.py b/tests/session/memory/test_memory_react.py index 9ae9d257fa..f3b9a3b69f 100644 --- a/tests/session/memory/test_memory_react.py +++ b/tests/session/memory/test_memory_react.py @@ -4,18 +4,19 @@ Tests for memory ExtractLoop orchestrator. """ +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest + from openviking.session.memory.dataclass import ( MemoryTypeSchema, ResolvedOperations, ) -from openviking.session.memory.schema_model_generator import SchemaModelGenerator - from openviking.session.memory.extract_loop import ( ExtractLoop, ) +from openviking.session.memory.schema_model_generator import SchemaModelGenerator class TestPreFetchFileFiltering: @@ -248,6 +249,93 @@ def test_final_skeleton_always_includes_delete_ids(self): "preferences": [], } + @pytest.mark.asyncio + async def test_pseudo_tool_json_retries_before_accepting_empty_operations(self, monkeypatch): + class FakeVLM: + model = "test-model" + + def __init__(self): + self.responses = iter( + [ + '{"tool_call_name":"search","args":{"query":"Melanie pottery"}}', + '{"preferences":[],"delete_ids":[]}', + ] + ) + self.seen_messages = [] + + async def get_completion_async(self, **kwargs): + self.seen_messages.append(list(kwargs["messages"])) + return next(self.responses) + + class FakeContextProvider: + read_file_contents = {} + + def get_memory_schemas(self, ctx): + return [ + MemoryTypeSchema( + memory_type="preferences", + description="Preferences", + directory="viking://user/{user_space}/memories/preferences", + filename_template="{topic}.md", + fields=[], + ) + ] + + def get_tools(self): + return [] + + def get_extract_context(self): + return MagicMock() + + def get_output_language(self): + return "en" + + def instruction(self): + return "Extract memory operations." + + async def prefetch(self): + return [] + + vlm = FakeVLM() + isolation_handler = MagicMock() + isolation_handler.get_read_scope.return_value = None + extract_loop = ExtractLoop( + vlm=vlm, + viking_fs=MagicMock(), + context_provider=FakeContextProvider(), + isolation_handler=isolation_handler, + max_iterations=1, + ) + config = SimpleNamespace(memory=SimpleNamespace(link_enabled=False)) + monkeypatch.setattr( + "openviking.session.memory.extract_loop.get_openviking_config", + lambda: config, + ) + monkeypatch.setattr( + "openviking_cli.utils.config.get_openviking_config", + lambda: config, + ) + resolved = ResolvedOperations( + upsert_operations=[], + delete_file_contents=[], + errors=[], + ) + extract_loop.resolve_operations = AsyncMock(return_value=(resolved, [])) + extract_loop._check_unread_existing_files = AsyncMock(return_value={}) + extract_loop._validate_patch_operations = MagicMock(return_value=[]) + extract_loop.finalize_operations = AsyncMock() + + operations, tools_used = await extract_loop.run() + + assert len(vlm.seen_messages) == 2 + assert operations.upsert_operations == [] + assert operations.delete_file_contents == [] + assert tools_used == [] + assert any( + "previous output could not be parsed as valid JSON" in message.get("content", "") + for message in vlm.seen_messages[1] + ) + @pytest.mark.asyncio async def test_final_unparseable_response_raises_instead_of_empty_success(self): class FakeVLM: diff --git a/tests/session/memory/test_memory_react_system_prompt.py b/tests/session/memory/test_memory_react_system_prompt.py index a37bdc02e6..4b0aef9a65 100644 --- a/tests/session/memory/test_memory_react_system_prompt.py +++ b/tests/session/memory/test_memory_react_system_prompt.py @@ -4,29 +4,58 @@ Test that provider instruction correctly instructs LLM. """ +from types import SimpleNamespace + from openviking.message import ImagePart, Message, TextPart, ToolPart from openviking.session.memory.session_extract_context_provider import SessionExtractContextProvider from openviking.session.memory.vision_message_normalizer import IMAGE_DESCRIPTION_PROMPT +def _provider_with_prefetch_mode(monkeypatch, *, eager_prefetch: bool): + config = SimpleNamespace( + memory=SimpleNamespace( + eager_prefetch=eager_prefetch, + prefetch_search_topn=5, + link_enabled=False, + ) + ) + monkeypatch.setattr( + "openviking.session.memory.session_extract_context_provider.get_openviking_config", + lambda: config, + ) + monkeypatch.setattr( + "openviking.session.memory.utils.resolve_output_language", + lambda _text: "en", + ) + return SessionExtractContextProvider(messages=[]) + + class TestProviderInstruction: """Test the provider instruction contains correct instructions.""" - def test_instruction_contains_read_before_edit_instructions(self): - """Test that instruction explicitly tells LLM to read files before editing.""" - # Create provider with mock messages - mock_messages = [] - provider = SessionExtractContextProvider(messages=mock_messages) + def test_eager_prefetch_instruction_says_no_tools_are_available(self, monkeypatch): + provider = _provider_with_prefetch_mode(monkeypatch, eager_prefetch=True) instruction = provider.instruction() - # Check for critical instructions + assert "No tools are available" in instruction + assert "available tools (read/search)" not in instruction + assert "ONLY read and search tools are available" not in instruction + + def test_on_demand_instruction_exposes_only_read_tool(self, monkeypatch): + """Non-eager extraction may read listed URIs, but cannot search or write.""" + provider = _provider_with_prefetch_mode(monkeypatch, eager_prefetch=False) + + instruction = provider.instruction() + + assert "ONLY the read tool is available" in instruction + assert "search and write are not available" in instruction assert ( "Before editing ANY existing memory file, you MUST first read its complete content" in instruction ) assert ( - "ONLY read URIs that are explicitly listed in ls/search tool results, returned by previous tool calls" + "ONLY read URIs that are explicitly listed in pre-fetched search results, returned by previous tool calls" in instruction ) @@ -232,8 +261,6 @@ def test_detect_language_prefers_user_text_over_assistant_text(self): assert provider._detect_language() == "zh-CN" - - async def test_prepare_extraction_messages_replaces_image_part_with_vlm_description(self): class FakeVisionVLM: def __init__(self): @@ -341,6 +368,7 @@ async def test_prepare_extraction_messages_does_not_replace_caller_message_list( assert any(isinstance(part, ImagePart) for part in messages[0].parts) assert provider.messages is not messages + def test_session_provider_empty_messages_still_uses_environment_fallback(monkeypatch): monkeypatch.setenv("TZ", "Asia/Shanghai") provider = SessionExtractContextProvider(messages=[]) diff --git a/tests/session/memory/test_memory_tools.py b/tests/session/memory/test_memory_tools.py index f0161c0946..0d79fd3dcc 100644 --- a/tests/session/memory/test_memory_tools.py +++ b/tests/session/memory/test_memory_tools.py @@ -4,6 +4,8 @@ Tests for memory tools. """ +import json + import pytest from openviking.server.identity import RequestContext, Role, ToolContext @@ -13,6 +15,7 @@ MemoryLsTool, MemoryReadTool, MemorySearchTool, + add_tool_call_pair_to_messages, get_tool, ) from openviking_cli.session.user_id import UserIdentifier @@ -21,6 +24,23 @@ class TestMemoryTools: """Tests for memory tools.""" + def test_prefetched_tool_result_is_not_serialized_as_a_tool_call(self): + messages = [] + + add_tool_call_pair_to_messages( + messages, + call_id="prefetch-1", + tool_name="search", + params={"query": "Melanie pottery"}, + result=["viking://user/default/memories/events/pottery.md"], + ) + + payload = json.loads(messages[0]["content"]) + assert messages[0]["role"] == "user" + assert payload["message_type"] == "tool_result" + assert payload["tool_name"] == "search" + assert "tool_call_name" not in payload + def test_read_tool_properties(self): """Test MemoryReadTool properties.""" tool = MemoryReadTool() diff --git a/tests/session/memory/test_patch_merge_context_provider.py b/tests/session/memory/test_patch_merge_context_provider.py index 39c98fdb28..e64c09faac 100644 --- a/tests/session/memory/test_patch_merge_context_provider.py +++ b/tests/session/memory/test_patch_merge_context_provider.py @@ -61,7 +61,9 @@ async def test_patch_merge_context_provider_prefetch_reads_originals_and_renders assert provider.get_tools() == [] assert provider.read_file.await_count == 1 read_message = json.loads(messages[0]["content"]) - assert read_message["tool_call_name"] == "read" + assert read_message["message_type"] == "tool_result" + assert read_message["tool_name"] == "read" + assert "tool_call_name" not in read_message assert read_message["args"] == {"uri": "viking://user/u/memories/experiences/booking.md"} assert read_message["result"]["experience_name"] == "booking" assert messages[1]["role"] == "user"