diff --git a/openviking/session/memory/memory_isolation_handler.py b/openviking/session/memory/memory_isolation_handler.py index f9f4fe1bbf..9695a8641e 100644 --- a/openviking/session/memory/memory_isolation_handler.py +++ b/openviking/session/memory/memory_isolation_handler.py @@ -143,6 +143,17 @@ def _unique_peer_target_id_in_messages(self) -> Optional[str]: peer_ids = list(dict.fromkeys(targets)) return peer_ids[0] if len(peer_ids) == 1 else None + def _unique_owner_target_id_in_messages(self) -> Optional[str]: + targets = [] + for msg in self._messages(): + if not self._is_peer_owner_message(msg): + continue + target_id = self._message_target_id(msg) + if target_id: + targets.append(target_id) + target_ids = list(dict.fromkeys(targets)) + return target_ids[0] if len(target_ids) == 1 else None + def render_schema_directories(self, memory_type_schema: MemoryTypeSchema) -> List[str]: user_id = self.ctx.user.user_id if self.ctx and self.ctx.user else "default" user_space = user_id @@ -174,12 +185,24 @@ def _range_targets(self, ranges: Any) -> List[str]: return [] target_ids = [] + has_messages = False + has_owner_message = False for msg_group in getattr(msg_range, "elements", []) or []: for msg in msg_group: + has_messages = True + has_owner_message = has_owner_message or self._is_peer_owner_message(msg) target_id = self._message_target_id(msg) if target_id: target_ids.append(target_id) - return list(dict.fromkeys(target_ids)) + target_ids = list(dict.fromkeys(target_ids)) + if target_ids: + return target_ids + + if not has_messages or has_owner_message: + return [] + + fallback_target_id = self._unique_owner_target_id_in_messages() + return [fallback_target_id] if fallback_target_id else [] def _resolve_operation_target_id(self, raw_peer_id: Any) -> Optional[str]: peer_id = safe_peer_id(raw_peer_id) diff --git a/tests/session/memory/test_memory_isolation_handler.py b/tests/session/memory/test_memory_isolation_handler.py index 21ccdc3aa7..ec4906b1a5 100644 --- a/tests/session/memory/test_memory_isolation_handler.py +++ b/tests/session/memory/test_memory_isolation_handler.py @@ -543,6 +543,126 @@ def test_calculate_memory_uris_ranges_ignore_assistant_peer_ids(self, mock_gener assert operation.memory_fields["user_id"] == "support_bot" assert "peer_id" not in operation.memory_fields + @patch("openviking.session.memory.memory_isolation_handler.generate_uri") + def test_calculate_memory_uris_assistant_only_range_uses_unique_session_target( + self, mock_generate_uri + ): + mock_generate_uri.side_effect = lambda **kwargs: ( + f"viking://user/{kwargs.get('user_space')}/memories/events/demo" + ) + + ctx = create_ctx(user_id="support_bot") + owner_message = create_message("user", "request") + ranged_message = create_message("assistant", "completed request", peer_id="assistant-bot") + messages = [owner_message, ranged_message] + extract_ctx = create_mock_extract_context(messages) + mock_range = MagicMock() + mock_range.elements = [[ranged_message]] + extract_ctx.read_message_ranges.return_value = mock_range + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=True, + allowed_peer_ids={"assistant-bot"}, + ) + + from openviking.session.memory.dataclass import MemoryTypeSchema, ResolvedOperation + + schema = MemoryTypeSchema( + memory_type="events", + filename_template="demo.md", + directory="viking://user/{user_space}/memories/events", + ) + operation = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"event_name": "demo", "ranges": "1"}, + memory_type="events", + uris=[], + ) + + uris = handler.calculate_memory_uris(schema, operation, extract_ctx) + + assert uris == ["viking://user/support_bot/memories/events/demo"] + assert "peer_id" not in operation.memory_fields + + @patch("openviking.session.memory.memory_isolation_handler.generate_uri") + def test_calculate_memory_uris_assistant_only_range_drops_ambiguous_session_target( + self, mock_generate_uri + ): + ctx = create_ctx(user_id="support_bot") + ranged_message = create_message("assistant", "completed request", peer_id="assistant-bot") + messages = [ + create_message("user", "self request"), + create_message("user", "peer request", peer_id="web-visitor-alice"), + ranged_message, + ] + extract_ctx = create_mock_extract_context(messages) + mock_range = MagicMock() + mock_range.elements = [[ranged_message]] + extract_ctx.read_message_ranges.return_value = mock_range + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=True, + allowed_peer_ids={"assistant-bot", "web-visitor-alice"}, + ) + + from openviking.session.memory.dataclass import MemoryTypeSchema, ResolvedOperation + + schema = MemoryTypeSchema( + memory_type="events", + filename_template="demo.md", + directory="viking://user/{user_space}/memories/events", + ) + operation = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"event_name": "demo", "ranges": "2"}, + memory_type="events", + uris=[], + ) + + uris = handler.calculate_memory_uris(schema, operation, extract_ctx) + + assert uris == [] + mock_generate_uri.assert_not_called() + + @patch("openviking.session.memory.memory_isolation_handler.generate_uri") + def test_calculate_memory_uris_unallowed_user_range_does_not_use_session_fallback( + self, mock_generate_uri + ): + ctx = create_ctx(user_id="support_bot") + owner_message = create_message("user", "self request") + ranged_message = create_message("user", "other peer", peer_id="web-visitor-bob") + extract_ctx = create_mock_extract_context([owner_message, ranged_message]) + mock_range = MagicMock() + mock_range.elements = [[ranged_message]] + extract_ctx.read_message_ranges.return_value = mock_range + handler = MemoryIsolationHandler( + ctx, + extract_ctx, + allow_self=True, + allowed_peer_ids={"web-visitor-alice"}, + ) + + from openviking.session.memory.dataclass import MemoryTypeSchema, ResolvedOperation + + schema = MemoryTypeSchema( + memory_type="events", + filename_template="demo.md", + directory="viking://user/{user_space}/memories/events", + ) + operation = ResolvedOperation( + old_memory_file_content=None, + memory_fields={"event_name": "demo", "ranges": "1"}, + memory_type="events", + uris=[], + ) + + uris = handler.calculate_memory_uris(schema, operation, extract_ctx) + + assert uris == [] + mock_generate_uri.assert_not_called() + @patch("openviking.session.memory.memory_isolation_handler.generate_uri") def test_calculate_memory_uris_unallowed_peer_id_does_not_fallback(self, mock_generate_uri): mock_generate_uri.side_effect = lambda **kwargs: (