Fix/crashing menus

src/jwinventoryapi/manager/container_manager/container_adapter.py

@@ -0,0 +1,52 @@
+from typing import Protocol, TYPE_CHECKING
+
+from endstone import Player
+from endstone.inventory import ItemStack
+
+if TYPE_CHECKING:
+    from jwinventoryapi.menu.menu_inventory import MenuInventory
+
+
+class ContainerAdapter(Protocol):
+    def get(self, slot: int) -> ItemStack | None: ...
+    def set(self, slot: int, item: ItemStack | None) -> None: ...
+
+
+class PlayerInventoryAdapter:
+    def __init__(self, player: Player):
+        self.player = player
+
+    def get(self, slot: int) -> ItemStack | None:
+        return self.player.inventory.get_item(slot)
+
+    def set(self, slot: int, item: ItemStack | None) -> None:
+        if item is None:
+            self.player.inventory.clear(slot)
+        else:
+            self.player.inventory.set_item(slot, item)
+
+
+class VirtualInventoryAdapter:
+    def __init__(self, inventory: "MenuInventory"):
+        self.inventory = inventory
+
+    def get(self, slot: int) -> ItemStack | None:
+        return self.inventory.get_item(slot)
+
+    def set(self, slot: int, item: ItemStack | None) -> None:
+        if item is None:
+            self.inventory.clear(slot)
+        else:
+            self.inventory.set_item(slot, item)
+
+
+class CursorAdapter:
+
+    def __init__(self):
+        self.cursor_item: ItemStack | None = None
+
+    def get(self, slot: int = 0) -> ItemStack | None:
+        return self.cursor_item
+
+    def set(self, slot: int, item: ItemStack | None) -> None:
+        self.cursor_item = item

src/jwinventoryapi/manager/container_manager/container_manager.py

@@ -0,0 +1,159 @@
+from typing import TYPE_CHECKING
+
+from bedrock_protocol.packets.types import ItemStackRequestSlotInfo
+from bedrock_protocol.packets.types.item_stack_response import ItemStackResponse
+from endstone import Player
+from endstone.inventory import ItemStack
+
+from jwinventoryapi.network.container_ui_ids import ContainerUIIds
+from jwinventoryapi.util.item_utils import pop_item, can_stack, clone_item, is_air
+
+from .container_adapter import ContainerAdapter, PlayerInventoryAdapter, VirtualInventoryAdapter, CursorAdapter
+from .transaction_container import TransactionContainer
+from .item_stack_response_builder import ItemStackResponseBuilder
+from .item_stack_tracker import ItemStackTracker
+
+if TYPE_CHECKING:
+    from jwinventoryapi.menu.menu_inventory import MenuInventory
+
+
+class ContainerManager:
+
+    def __init__(self, player: Player, inventory: "MenuInventory"):
+        self.player = player
+        self.tracker = ItemStackTracker()
+        self.player_container = TransactionContainer(PlayerInventoryAdapter(player))
+        self.virtual_container = TransactionContainer(VirtualInventoryAdapter(inventory))
+        self.cursor_container = TransactionContainer(CursorAdapter())
+        self._response_builder: ItemStackResponseBuilder | None = None
+        self._modified_player_slots: set[int] = set()
+
+    def begin_request(self, request_id: int) -> None:
+        self._response_builder = ItemStackResponseBuilder(request_id, self.tracker, self)
+
+    @staticmethod
+    def _is_player_container(container_enum: int) -> bool:
+        return container_enum in (
+            ContainerUIIds.INVENTORY,
+            ContainerUIIds.HOTBAR,
+            ContainerUIIds.COMBINED_HOTBAR_AND_INVENTORY,
+        )
+
+    def track_request_slot(self, slot_info: ItemStackRequestSlotInfo) -> None:
+        self.tracker.seed_from_request(slot_info)
+        if self._is_player_container(slot_info.container.container_enum):
+            self._modified_player_slots.add(slot_info.slot)
+        if self._response_builder is not None:
+            self._response_builder.add_slot(slot_info.container.container_enum, slot_info.slot)
+
+    def get_container_adapter_and_slot(self, slot_info: ItemStackRequestSlotInfo) -> tuple[ContainerAdapter, int] | None:
+        container_type = slot_info.container.container_enum
+        slot = slot_info.slot
+        if container_type == ContainerUIIds.LEVEL_ENTITY:
+            return self.virtual_container, slot
+        elif container_type == ContainerUIIds.CURSOR:
+            return self.cursor_container, slot
+        elif container_type in (ContainerUIIds.INVENTORY, ContainerUIIds.HOTBAR, ContainerUIIds.COMBINED_HOTBAR_AND_INVENTORY):
+            return self.player_container, slot
+        raise ValueError(f"Unsupported container type: {container_type}")
+
+    def get_item_at(self, container_enum: int, slot: int) -> ItemStack | None:
+        if container_enum == ContainerUIIds.LEVEL_ENTITY:
+            return self.virtual_container.actual.get(slot)
+        if container_enum == ContainerUIIds.CURSOR:
+            return self.cursor_container.actual.get(slot)
+        if container_enum in (ContainerUIIds.INVENTORY, ContainerUIIds.HOTBAR, ContainerUIIds.COMBINED_HOTBAR_AND_INVENTORY):
+            return self.player_container.actual.get(slot)
+        raise ValueError(f"Unsupported container type: {container_enum}")
+
+    def assign_virtual_slot(self, slot: int, item: ItemStack | None) -> int:
+        return self.tracker.assign_slot(ContainerUIIds.LEVEL_ENTITY, slot, item)
+
+    def transfer_items(self, source: ItemStackRequestSlotInfo, destination: ItemStackRequestSlotInfo, count: int):
+        self.track_request_slot(source)
+        self.track_request_slot(destination)
+        removed = self.remove_item_from_slot(source, count)
+        self.add_item_to_slot(destination, removed, count)
+
+    def remove_item_from_slot(self, slot_info: ItemStackRequestSlotInfo, count: int) -> ItemStack:
+        self.track_request_slot(slot_info)
+        container, slot = self.get_container_adapter_and_slot(slot_info)
+        if count < 1:
+            raise ValueError("item count is less than 1")
+        item: ItemStack = container.get(slot)
+        if item is None or item.amount < count:
+            raise ValueError("item is None or existing item < count")
+        removed, remainder = pop_item(item, count)
+        if remainder is None or is_air(remainder):
+            container.set(slot, None)
+        else:
+            container.set(slot, remainder)
+        return removed
+
+    def add_item_to_slot(self, slot_info: ItemStackRequestSlotInfo, item: ItemStack, count: int):
+        self.track_request_slot(slot_info)
+        if item is None or is_air(item):
+            raise ValueError("item is None or existing item < count")
+        container, slot = self.get_container_adapter_and_slot(slot_info)
+        if count < 1:
+            return
+        existing_item: ItemStack = container.get(slot) or ItemStack("minecraft:air")
+        if existing_item is None or is_air(existing_item):
+            container.set(slot, clone_item(item))
+            return
+        if not can_stack(existing_item, item):
+            raise ValueError("cannot stack items")
+        merged = clone_item(existing_item)
+        merged.amount += item.amount
+        container.set(slot, merged)
+
+    def handle_swap(self, slot1: ItemStackRequestSlotInfo, slot2: ItemStackRequestSlotInfo):
+        self.track_request_slot(slot1)
+        self.track_request_slot(slot2)
+        c1, s1 = self.get_container_adapter_and_slot(slot1)
+        c2, s2 = self.get_container_adapter_and_slot(slot2)
+        item1 = c1.get(s1)
+        item2 = c2.get(s2)
+        c1.set(s1, item2)
+        c2.set(s2, item1)
+
+    def handle_drop(self, source: ItemStackRequestSlotInfo, count: int):
+        self.track_request_slot(source)
+        dropped = self.remove_item_from_slot(source, count)
+        if is_air(dropped):
+            raise ValueError("cannot drop empty item stack")
+        if dropped.amount > dropped.max_stack_size:
+            raise ValueError("cannot drop item stack larger than max stack size")
+        self.player.dimension.drop_item(self.player.location, dropped)
+
+    def commit_transaction(self) -> ItemStackResponse:
+        self.player_container.commit()
+        self.virtual_container.commit()
+        self.cursor_container.commit()
+        if self._response_builder is None:
+            raise RuntimeError("No active item stack request")
+        for container_enum, slots in self._response_builder.changed_slots().items():
+            for slot in slots.values():
+                item = self.get_item_at(container_enum, slot)
+                self.tracker.assign_slot(container_enum, slot, item)
+        response = self._response_builder.build()
+        self._response_builder = None
+        return response
+
+    def discard_transaction(self) -> None:
+        self.player_container.discard()
+        self.virtual_container.discard()
+        self.cursor_container.discard()
+        self._response_builder = None
+
+    def sync_player_inventory(self) -> None:
+        if not self._modified_player_slots:
+            return
+        inventory = self.player.inventory
+        for slot in self._modified_player_slots:
+            item = inventory.get_item(slot)
+            inventory.clear(slot)
+            if item is not None and not is_air(item):
+                inventory.set_item(slot, clone_item(item))
+        self._modified_player_slots.clear()
+        self.tracker.clear_player_slots()

src/jwinventoryapi/manager/container_manager/item_stack_response_builder.py

@@ -0,0 +1,79 @@
+from typing import TYPE_CHECKING
+
+from bedrock_protocol.packets.types import FullContainerName
+from bedrock_protocol.packets.types.item_stack_response import (
+    ItemStackResponse,
+    ItemStackResponseContainerInfo,
+    ItemStackResponseSlotInfo,
+)
+from endstone.inventory import ItemStack
+
+from jwinventoryapi.util.item_utils import is_air
+
+from .item_stack_tracker import ItemStackTracker
+
+if TYPE_CHECKING:
+    from .container_manager import ContainerManager
+
+
+class ItemStackResponseBuilder:
+    def __init__(self, request_id: int, tracker: ItemStackTracker, container_manager: "ContainerManager"):
+        self._request_id = request_id
+        self._tracker = tracker
+        self._container_manager = container_manager
+        self._changed_slots: dict[int, dict[int, int]] = {}
+
+    def add_slot(self, container_enum: int, slot: int) -> None:
+        self._changed_slots.setdefault(container_enum, {})[slot] = slot
+
+    @staticmethod
+    def _get_custom_name(item: ItemStack | None) -> str:
+        if item is None or is_air(item):
+            return ""
+        item_meta = item.item_meta
+        if item_meta is not None and item_meta.has_display_name:
+            return item_meta.display_name
+        return ""
+
+    @staticmethod
+    def _get_durability_correction(item: ItemStack | None) -> int:
+        if item is None or is_air(item):
+            return 0
+        return item.data
+
+    def build(self) -> ItemStackResponse:
+        container_infos: list[ItemStackResponseContainerInfo] = []
+        for container_enum, slots in self._changed_slots.items():
+            slot_infos: list[ItemStackResponseSlotInfo] = []
+            for slot in slots.values():
+                item = self._container_manager.get_item_at(container_enum, slot)
+                stack_id = self._tracker.get_stack_id(container_enum, slot)
+                custom_name = self._get_custom_name(item)
+                slot_infos.append(ItemStackResponseSlotInfo(
+                    slot=slot,
+                    hotbar_slot=slot,
+                    count=0 if item is None or is_air(item) else item.amount,
+                    item_stack_id=stack_id,
+                    custom_name=custom_name,
+                    filtered_custom_name=custom_name,
+                    durability_correction=self._get_durability_correction(item),
+                ))
+            container_infos.append(ItemStackResponseContainerInfo(
+                container=FullContainerName(container_enum),
+                slots=slot_infos,
+            ))
+        return ItemStackResponse(
+            result=ItemStackResponse.RESULT_OK,
+            request_id=self._request_id,
+            container_infos=container_infos,
+        )
+
+    def changed_slots(self) -> dict[int, dict[int, int]]:
+        return self._changed_slots
+
+    @classmethod
+    def build_error(cls, request_id: int) -> ItemStackResponse:
+        return ItemStackResponse(
+            result=ItemStackResponse.RESULT_ERROR,
+            request_id=request_id,
+        )

src/jwinventoryapi/manager/container_manager/item_stack_tracker.py

@@ -0,0 +1,43 @@
+from bedrock_protocol.packets.types import ItemStackRequestSlotInfo
+from endstone.inventory import ItemStack
+
+from jwinventoryapi.network.container_ui_ids import ContainerUIIds
+from jwinventoryapi.util.item_utils import is_air
+
+
+class ItemStackTracker:
+    def __init__(self):
+        self._next_stack_id = 1
+        self._stack_ids: dict[tuple[int, int], int] = {}
+
+    def seed_from_request(self, slot_info: ItemStackRequestSlotInfo) -> None:
+        key = (slot_info.container.container_enum, slot_info.slot)
+        if key in self._stack_ids:
+            return
+        if slot_info.net_id > 0:
+            self._stack_ids[key] = slot_info.net_id
+            self._next_stack_id = max(self._next_stack_id, slot_info.net_id + 1)
+
+    def assign_slot(self, container_enum: int, slot: int, item: ItemStack | None) -> int:
+        if item is None or is_air(item):
+            stack_id = 0
+        else:
+            stack_id = self._next_stack_id
+            self._next_stack_id += 1
+        self._stack_ids[(container_enum, slot)] = stack_id
+        return stack_id
+
+    def get_stack_id(self, container_enum: int, slot: int) -> int:
+        return self._stack_ids.get((container_enum, slot), 0)
+
+    def clear_player_slots(self) -> None:
+        player_containers = (
+            ContainerUIIds.INVENTORY,
+            ContainerUIIds.HOTBAR,
+            ContainerUIIds.COMBINED_HOTBAR_AND_INVENTORY,
+        )
+        self._stack_ids = {
+            key: stack_id
+            for key, stack_id in self._stack_ids.items()
+            if key[0] not in player_containers
+        }

src/jwinventoryapi/manager/container_manager/transaction_container.py

@@ -0,0 +1,33 @@
+from endstone.inventory import ItemStack
+
+from jwinventoryapi.manager.container.container_adapter import ContainerAdapter
+from jwinventoryapi.util.item_utils import is_air, clone_item
+
+
+class TransactionContainer:
+    def __init__(self, actual: ContainerAdapter):
+        self.actual = actual
+        self.changed_slots: dict[int, ItemStack | None] = {}
+
+    def get(self, slot: int) -> ItemStack | None:
+        if slot in self.changed_slots:
+            item = self.changed_slots[slot]
+        else:
+            item = self.actual.get(slot)
+        if item is None or is_air(item):
+            return item
+        return clone_item(item)
+
+    def set(self, slot: int, item: ItemStack | None) -> None:
+        if item is None or is_air(item):
+            self.changed_slots[slot] = None
+        else:
+            self.changed_slots[slot] = clone_item(item)
+
+    def commit(self) -> None:
+        for slot, item in self.changed_slots.items():
+            self.actual.set(slot, item)
+        self.changed_slots.clear()
+
+    def discard(self) -> None:
+        self.changed_slots.clear()

src/jwinventoryapi/manager/player_manager.py

@@ -20,9 +20,9 @@ def create_session(player: Player) -> 'Session':
     return session
 
 
-def close_session(player: Player):
+def close_session(player: Player, sync_inventory: bool = False):
     session = sessions.pop(player.unique_id, None)
     if session is not None:
         if session.menu is not None:
             session.menu._remove_session(session)
-        session.close()
+        session.close(sync_inventory=sync_inventory)