"""Admin writes must be serialized with runtime writes of the same installation. Review F1: ``update_settings`` / ``set_session_limit`` / ``reset_session`` / ``reset_all`` / ``delete_session`` mutated the cached state and then awaited ``persist`` without taking the installation lock that ``check_and_count`` holds. ``persist`` overwrites the whole row, so when a Host completes two in-flight writes out of call order, the older snapshot lands last and silently rolls the newer one back. The synthetic Host below keeps ``SET_PLUGIN_STORAGE`` pending until the test releases it, and releases newest-first, which is the order that makes an unserialized pair lose the newer write. Pre-fix, an admin settings write and a counter write of one installation both sit in the Host queue at once, and the durable row ends up missing the session the counter write had recorded. """ from __future__ import annotations import asyncio import base64 import importlib.util import json import sys from pathlib import Path import pytest from langbot_plugin.api.proxies.invocation import bind_invocation from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction as Action from langbot_plugin.entities.io.context import InstallationBinding PLUGIN_ROOT = Path(__file__).resolve().parents[1] def _load_plugin_module(module_name: str): """Load this plugin's ``main.py`` under a unique name.""" spec = importlib.util.spec_from_file_location(module_name, PLUGIN_ROOT / "main.py") module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module spec.loader.exec_module(module) return module plugin_main = _load_plugin_module("daily_limit_plugin_serialization_under_test") DailyLimitPlugin = plugin_main.DailyLimitPlugin INSTANCE = "instance-1" WORKSPACE = "workspace-1" CONFIG = { "daily_limit": 50, "limit_message": "configured message", "silent_mode": False, "reset_timezone_offset": 8, "reset_hour": 0, } class ReorderingHost: """Host KV whose writes complete only when the test releases them.""" def __init__(self) -> None: self.rows: dict[tuple[str, str, str], bytes] = {} self.pending: list[tuple[str, bytes, asyncio.Future]] = [] self.commits: list[str] = [] async def call_action(self, action, data): key = data.get("key", "") row = (INSTANCE, WORKSPACE, key) if action == Action.GET_PLUGIN_STORAGE: return {"value_base64": base64.b64encode(self.rows.get(row, b"")).decode()} if action == Action.SET_PLUGIN_STORAGE: value = base64.b64decode(data["value_base64"]) future = asyncio.get_running_loop().create_future() self.pending.append((key, value, future)) await future return {} if action == Action.DELETE_PLUGIN_STORAGE: self.rows.pop(row, None) return {} raise AssertionError(action) def release_newest(self) -> bool: """Commit the most recently issued pending write, as a busy Host can.""" if not self.pending: return False key, value, future = self.pending.pop() self.rows[(INSTANCE, WORKSPACE, key)] = value self.commits.append(key) future.set_result(None) return True class FlakyHost(ReorderingHost): """Host that fails exactly one write, the ``fail_on``-th, and accepts the rest.""" def __init__(self, fail_on: int) -> None: super().__init__() self.fail_on = fail_on self.writes_seen = 0 async def call_action(self, action, data): if action == Action.SET_PLUGIN_STORAGE: self.writes_seen += 1 if self.writes_seen == self.fail_on: raise RuntimeError("host storage unavailable") return await super().call_action(action, data) def _binding(installation: str) -> InstallationBinding: return InstallationBinding( instance_uuid=INSTANCE, workspace_uuid=WORKSPACE, installation_uuid=installation, runtime_revision=1, artifact_digest="a" * 64, ) def _plugin(host: ReorderingHost) -> DailyLimitPlugin: plugin = DailyLimitPlugin() plugin.plugin_runtime_handler = host return plugin async def _settle(turns: int = 20) -> None: for _ in range(turns): await asyncio.sleep(0) async def _drain(host: ReorderingHost, *tasks: asyncio.Task) -> None: """Run the tasks to completion while the Host commits writes newest-first.""" for _ in range(200): if host.release_newest(): await _settle() continue if all(task.done() for task in tasks): break await asyncio.sleep(0) assert all(task.done() for task in tasks), "the Host writes never completed" for task in tasks: task.result() def _row(host: ReorderingHost, binding: InstallationBinding) -> dict: scope = plugin_main.installation_scope(binding) return json.loads(host.rows[(INSTANCE, WORKSPACE, plugin_main.storage_key(scope))].decode("utf-8")) def test_runtime_write_is_not_overwritten_by_an_overlapping_admin_write(): """A counter write and an admin settings write of one installation must not interleave.""" async def scenario(): host = ReorderingHost() plugin = _plugin(host) binding = _binding("installation-a") with bind_invocation(host, config=CONFIG, binding=binding): settings_write = asyncio.create_task(plugin.update_settings({"limit_message": "admin"})) await _settle() # PRE-FIX: the admin write reaches the Host without the installation # lock, so the runtime write below starts while it is still in flight. assert len(host.pending) == 1, "the settings write never reached the Host" counter_write = asyncio.create_task(plugin.check_and_count("group:1", "group:1")) await _settle(50) await _drain(host, settings_write, counter_write) row = _row(host, binding) # PRE-FIX: the older settings snapshot was committed last, so the durable # row had no "group:1" session at all. assert list(row["sessions"]) == ["group:1"] assert row["sessions"]["group:1"]["count"] == 1 assert row["settings"]["limit_message"] == "admin" # The durable row is the truth: a reload must show the same state. reloaded = _plugin(host) with bind_invocation(host, config=CONFIG, binding=binding): snapshot = await reloaded.snapshot() assert [row_["id"] for row_ in snapshot["sessions"]] == ["group:1"] assert snapshot["sessions"][0]["count"] == 1 assert snapshot["settings"]["limit_message"] == "admin" asyncio.run(scenario()) def test_admin_writes_of_one_installation_are_serialized_too(): """Two overlapping admin writes must not let the older snapshot win.""" async def scenario(): host = ReorderingHost() plugin = _plugin(host) binding = _binding("installation-a") with bind_invocation(host, config=CONFIG, binding=binding): counter_write = asyncio.create_task(plugin.check_and_count("group:1", "group:1")) await _drain(host, counter_write) assert list(_row(host, binding)["sessions"]) == ["group:1"] settings_write = asyncio.create_task(plugin.update_settings({"reset_hour": 5})) await _settle() assert len(host.pending) == 1, "the settings write never reached the Host" delete_write = asyncio.create_task(plugin.delete_session("group:1")) await _settle(50) await _drain(host, settings_write, delete_write) row = _row(host, binding) # PRE-FIX: the settings snapshot (which still held "group:1") was committed # after the deletion, so the session came back. assert row["sessions"] == {} assert row["settings"]["reset_hour"] == 5 asyncio.run(scenario()) def test_failed_write_does_not_leave_the_change_in_memory(): """A write that never landed must not stay in memory and leak into the next one.""" async def scenario(): host = FlakyHost(fail_on=2) plugin = _plugin(host) binding = _binding("installation-a") with bind_invocation(host, config=CONFIG, binding=binding): counter_write = asyncio.create_task(plugin.check_and_count("group:1", "group:1")) await _drain(host, counter_write) # PRE-FIX: the rejected value stayed in the cached settings dict and was # written out by the next successful write. with pytest.raises(RuntimeError): await plugin.update_settings({"limit_message": "never persisted"}) assert (await plugin.snapshot())["settings"]["limit_message"] == "configured message" await _drain(host, asyncio.create_task(plugin.reset_all())) row = _row(host, binding) assert row["settings"]["limit_message"] == "configured message" assert list(row["sessions"]) == ["group:1"] asyncio.run(scenario())