diff --git a/PyMemoryEditor/mcp/config.py b/PyMemoryEditor/mcp/config.py index 536a479..27ee114 100644 --- a/PyMemoryEditor/mcp/config.py +++ b/PyMemoryEditor/mcp/config.py @@ -46,7 +46,28 @@ #: value (``int`` ``0``, ``100``) legitimately matches millions of addresses; #: keeping them all would blow out the server's memory for a result set no #: refine loop can use anyway. At the cap the scan stops early and says so. -DEFAULT_MAX_SCAN_RESULTS = 50_000 +#: +#: 100 000 costs 4.2 MB per result set against 2.1 MB, and 21 ms to sort +#: against 10 ms. The figure that matters is not per set, though: a session +#: keeps ``MAX_SCANS_PER_SESSION`` (20) of them and a server keeps +#: ``MAX_OPEN_SESSIONS`` (8) sessions, so the ceiling on retained addresses +#: goes from ~336 MB to ~672 MB. Reaching it needs 160 capped result sets, +#: which a long refine chain across several targets can do. +#: +#: What it buys: a scan whose true hit count falls between the two ceilings +#: stops being flagged ``partial``, and refining a truncated set can converge +#: on an address that was never in it. That case is real but narrow — above +#: the new ceiling nothing changes, and a common value like ``int 0`` matches +#: millions and is out of reach at any sane cap. +#: +#: What it costs in time depends on the value's density, and an earlier +#: version of this comment got that wrong. Measured on ``int 0``, reaching +#: either ceiling took 0.02s against 0.03s — but that is the dense case, where +#: hits arrive faster than the clock can spend. For a sparser value the scan +#: has to walk further to collect twice as many hits, so the time roughly +#: doubles up to the 30-second budget. "Wall clock is unchanged" was true of +#: one measurement, not of the change. +DEFAULT_MAX_SCAN_RESULTS = 100_000 #: Wall-clock budget for one scan, in seconds. A full address-space scan of a #: large process takes minutes — long past the point where an MCP client gives diff --git a/PyMemoryEditor/mcp/server.py b/PyMemoryEditor/mcp/server.py index 26df82e..3b6ddfb 100644 --- a/PyMemoryEditor/mcp/server.py +++ b/PyMemoryEditor/mcp/server.py @@ -238,6 +238,22 @@ async def open_process(pid: int = 0, name: str = "", *, ctx: Context) -> Any: if not _can_elicit(ctx): raise SdkToolError(decision.reason) + # Before the prompt, not after. The cap lives in SessionStore.open, + # which runs once the handle is already open -- so without this the + # user was asked to approve an attach, approved it, and then got told + # the server was full. An approval is the most expensive step here and + # the one thing this server must not spend carelessly. + # + # The store renders the message, rather than this having its own: the + # first version here named neither the open sessions nor their scan + # counts, so the path a model actually hits was the less useful one. + refusal = toolset.store.capacity_refusal(found_pid) + if refusal is not None: + raise SdkToolError( + 'Not asking to attach to "%s" (pid %d): %s' + % (found_name or "unknown", found_pid, refusal) + ) + try: answer = await ctx.elicit( message=( diff --git a/PyMemoryEditor/mcp/session.py b/PyMemoryEditor/mcp/session.py index 684160a..ed80846 100644 --- a/PyMemoryEditor/mcp/session.py +++ b/PyMemoryEditor/mcp/session.py @@ -32,12 +32,19 @@ from typing import Dict, List, Optional, Sequence, Tuple from ..process.abstract import AbstractProcess +from ..process.util import pid_exists from ..process.region import MemoryRegion, MemoryRegionSnapshot #: Result sets retained per session before the oldest is evicted. MAX_SCANS_PER_SESSION = 20 +#: Processes one server keeps open at once. Refuses rather than evicting, +#: unlike the scan cap above: a session owns an OS handle, and closing one out +#: from under the model is not something it can detect. Bounds handle +#: exhaustion, since nothing obliges a model to call ``close_process``. +MAX_OPEN_SESSIONS = 8 + class SessionError(Exception): """An unknown or expired session/scan id. @@ -162,6 +169,22 @@ def scan_ids(self) -> Tuple[str, ...]: return tuple(self._scan_order) +def looks_alive(pid: int) -> bool: + """``pid_exists``, but a pid the OS will not even accept reads as dead. + + The reaper runs inside ``open``, so anything raised here escapes as + something other than a ``SessionError`` and the model is told only + "Error executing tool". ``pid_exists(2**31)`` raises ``OverflowError`` + (the value does not fit the platform's pid type), which the argument + fuzzer found by generating exactly that. A number the OS refuses cannot + name a running process, so treating it as dead is both safe and correct. + """ + try: + return pid_exists(pid) + except (OverflowError, ValueError, OSError): + return False + + class SessionStore: """The server's process handles and their scan results. @@ -178,14 +201,146 @@ def __init__(self) -> None: self._lock = threading.Lock() def open(self, process: AbstractProcess, pid: int, name: str) -> Session: - """Register an already-opened process and return its session.""" + """Register an already-opened process and return its session. + + :raises SessionError: if :data:`MAX_OPEN_SESSIONS` are already open. + The caller owns the handle it passed in and must close it. + """ + # Loop, rather than deciding to reap from a check taken earlier under + # a different acquisition: the store can fill between the two, and then + # a slot held by an exited process would be refused without anyone + # having looked. At most one reap — a second pass finds nothing new, + # and refusing has to terminate. + reaped = False + + while True: + with self._lock: + refusal = self._refusal_locked(pid) + + if refusal is None: + session_id = "proc-%d" % next(self._session_ids) + session = Session( + session_id=session_id, process=process, + pid=pid, name=name, + ) + self._sessions[session_id] = session + return session + + if reaped: + # Raised here, not after the block: computing the refusal + # under the lock and raising it outside leaves a window + # where another thread closes a session, and this caller + # is turned away with a message that is already false. + raise SessionError(refusal) + + # Outside the lock: reaping closes handles, which takes it again. + self._reap_dead() + reaped = True + + def _reap_dead(self) -> List[str]: + """Drop sessions whose target no longer exists. Caller holds no lock. + + The cap refuses rather than evicting, because closing a live handle + out from under the model is something it cannot detect. A *dead* + target is the exception that proves the rule: the handle is already + useless, so dropping it takes nothing away — and without this, eight + exited processes hold every slot forever and `open_process` is refused + for the rest of the server's life, with no way for the model to find + out which sessions are the dead ones (`process_info` answers happily + from cached state). + + A recycled pid reads as alive, which is inherent to asking the OS by + number. + """ + with self._lock: + dead = [ + session_id + for session_id, session in self._sessions.items() + if not looks_alive(session.pid) + ] + + for session_id in dead: + try: + self.close(session_id) + except SessionError: # closed concurrently + pass + return dead + + def capacity_refusal(self, pid: int = 0) -> Optional[str]: + """The reason :meth:`open` would refuse right now, or ``None``. + + One message, rendered once. The protocol layer asks this *before* + prompting the user — spending an approval on an attach that cannot + happen is worse than refusing outright — and :meth:`open` raises it as + the authoritative guard. Duplicating the text gave the early path a + worse message than the late one, which is the path a model actually + hits. + + Reaps first, or a server full of exited processes reports as full when + :meth:`open` would have made room. + """ + with self._lock: + if len(self._sessions) < MAX_OPEN_SESSIONS: + return None + + self._reap_dead() + with self._lock: - session_id = "proc-%d" % next(self._session_ids) - session = Session( - session_id=session_id, process=process, pid=pid, name=name + return self._refusal_locked(pid) + + def _refusal_locked(self, pid: int) -> Optional[str]: + """Render the refusal, or ``None``. **Caller holds** ``self._lock``. + + Separate from :meth:`capacity_refusal` so :meth:`open` can check and + insert under one acquisition. When the check released the lock and the + insert took it again, the cap could be exceeded: every thread saw room, + then every thread inserted. Demonstrated with 12 concurrent opens + against a cap of 8. + """ + # Deterministic, so a refactor that checks capacity outside the lock + # fails here instead of in a timing test. That regression happened + # once: the check and the insert were split across two acquisitions, + # every racing thread saw room, and the cap was exceeded. + assert self._lock.locked(), "capacity must be judged under the lock" + + if len(self._sessions) < MAX_OPEN_SESSIONS: + return None + + # No "close the oldest": the oldest is usually the target the whole + # session has been refining, and close_process discards its scan + # results too. The list is given instead, so the choice is made on + # what each session holds. + existing = [ + sid for sid, session in self._sessions.items() + if session.pid == pid + ] + already = ( + " Note that pid %d is already open as %s — reuse that id " + "rather than attaching again." % (pid, existing[0]) + if existing + else "" + ) + return ( + "This server already has %d processes open, which is the " + "limit, and all of them are live. Close whichever you are " + "done with (close_process also discards that session's scan " + "results) and attach again. Open: %s.%s" + % ( + MAX_OPEN_SESSIONS, + ", ".join( + "%s (%s, pid %d, %d scan%s)" + % (sid, session.name or "?", session.pid, + len(session.scan_ids), + "" if len(session.scan_ids) == 1 else "s") + for sid, session in self._sessions.items() + ), + already, ) - self._sessions[session_id] = session - return session + ) + + def at_capacity(self) -> bool: + """Whether :meth:`open` would refuse right now.""" + return self.capacity_refusal() is not None def get(self, session_id: str) -> Session: """Look up a session, or explain how to obtain a valid id.""" @@ -377,6 +532,8 @@ def host_platform() -> str: __all__ = ( + "MAX_OPEN_SESSIONS", + "looks_alive", "MAX_SCANS_PER_SESSION", "ScanResult", "Session", diff --git a/PyMemoryEditor/mcp/toolset.py b/PyMemoryEditor/mcp/toolset.py index c5f6c19..e7facc0 100644 --- a/PyMemoryEditor/mcp/toolset.py +++ b/PyMemoryEditor/mcp/toolset.py @@ -48,11 +48,15 @@ ) from .config import ServerConfig from .session import ( + MAX_OPEN_SESSIONS, + MAX_SCANS_PER_SESSION, ScanResult, Session, + SessionError, SessionStore, batch_regions, host_platform, + looks_alive, region_to_dict, ) @@ -553,12 +557,23 @@ def server_info(self) -> Dict[str, Any]: reachable, how long a scan may run, and which sessions are already open from earlier in the conversation. """ + # `alive` because a session outlives its target: nothing here notices + # a process exiting, `process_info` keeps answering from cached state, + # and the open-session cap means dead ones would hold slots the model + # cannot identify. The store reaps them when it needs room; this is how + # the model sees them before that. sessions = [ { "session_id": session.session_id, "pid": session.pid, "name": session.name, "scan_ids": list(session.scan_ids), + # `looks_alive`, not `pid_exists`: the raw call raises + # OverflowError for a pid outside the platform's range, and + # server_info is the tool the instructions tell the model to + # call first — a crash here hides the limits and the policy + # too, not just this field. + "alive": looks_alive(session.pid), } for session in self.store.sessions ] @@ -599,6 +614,8 @@ def server_info(self) -> Dict[str, Any]: # to. "max_address": format_address(MAX_ADDRESS), "max_offset_magnitude": format_address(MAX_OFFSET_MAGNITUDE), + "max_open_sessions": MAX_OPEN_SESSIONS, + "max_scans_per_session": MAX_SCANS_PER_SESSION, }, "open_sessions": sessions, } @@ -771,7 +788,17 @@ def attach(self, pid: int, name: str) -> Dict[str, Any]: "Could not open pid %d: %s. %s" % (pid, error, _permission_hint()) ) from None - session = self.store.open(process, pid, name) + try: + session = self.store.open(process, pid, name) + except SessionError: + # The handle already exists and the store's cap is checked after + # it does, so refusing without closing leaks the resource the cap + # protects. + try: + process.close() + except Exception: # noqa: BLE001 — target may already be gone + pass + raise result: Dict[str, Any] = { "opened": True, @@ -797,6 +824,10 @@ def open_process(self, pid: int = 0, name: str = "") -> Dict[str, Any]: reopening the same target, since each open costs a handle and resets the cached region map. + At most ``max_open_sessions`` (see ``server_info``) can be open at + once; reaching it is refused, not rotated, so ``close_process`` a + target you are done with. + Attaching to a process the operator has not pre-approved requires the **user's** approval, asked for at the moment you call this. Say which process you want and why; if the request is refused, report that rather @@ -2001,7 +2032,8 @@ def _read_addresses( """Yield ``(address, value | None)`` for each address, in one pass. Uses ``search_by_addresses``, which groups the reads by region so a - 50 000-address refine is a few hundred syscalls rather than 50 000. + refine of a capped result set is a few hundred syscalls rather than one + per address. """ yield from session.process.search_by_addresses( pytype, diff --git a/docs/mcp.md b/docs/mcp.md index b02896c..b8d025f 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -111,7 +111,7 @@ pymemoryeditor-mcp --allow-any-process # never asks at all (scripts/CI) --max-scan-results N - Addresses one scan may keep. Default 50 000. + Addresses one scan may keep. Default 100 000. --max-scan-seconds S @@ -129,7 +129,7 @@ pymemoryeditor-mcp --allow-any-process # never asks at all (scripts/CI) ToolWhat it does server_infoCapabilities, limits, policy and open sessions. The assistant should call this first. list_processesRunning processes the server is allowed to open. -open_processAttach by pid or name → a session_id. +open_processAttach by pid or name → a session_id. Eight targets can be open at once; past that it refuses rather than rotating one out, since a handle the assistant still holds must not vanish under it. close_processDetach and drop that session's scan results. process_infoBitness, address-space summary, modules and threads. list_memory_regionsPage through the memory map, filtered by permission or backing file. diff --git a/tests/mcp_server/test_parsing.py b/tests/mcp_server/test_parsing.py index c14ccad..cd42703 100644 --- a/tests/mcp_server/test_parsing.py +++ b/tests/mcp_server/test_parsing.py @@ -630,6 +630,55 @@ def test_the_partition_actually_covers_the_library(self): assert len(lib) + len(mcp) > len(sources) // 2 +class TestTheDocsQuoteTheRealLimits: + """`docs/mcp.md` spells the caps out as words, which a reader wants and a + constant change silently invalidates. Same guard as + `test_advertised_widths_are_exactly_the_accepted_widths`.""" + + WORDS = { + 1: "one", 2: "two", 3: "three", 4: "four", 5: "five", + 6: "six", 7: "seven", 8: "eight", 9: "nine", 10: "ten", + } + + def _mcp_doc(self): + from pathlib import Path + + return ( + Path(__file__).resolve().parents[2] / "docs" / "mcp.md" + ).read_text(encoding="utf-8") + + def test_the_open_session_cap_matches_the_constant(self): + from PyMemoryEditor.mcp.session import MAX_OPEN_SESSIONS + + # Looked up, not indexed. This test exists to turn a silent doc/code + # divergence into an actionable message, and `WORDS[12]` would have + # replaced that with `KeyError: 12` — loud, but saying nothing about + # what to do. + word = self.WORDS.get(MAX_OPEN_SESSIONS) + assert word is not None, ( + "MAX_OPEN_SESSIONS is %d and WORDS has no spelling for it. Add " + "one, then update the open_process row in docs/mcp.md to match." + % MAX_OPEN_SESSIONS + ) + + assert "%s targets can be open at once" % word in self._mcp_doc().lower(), ( + "docs/mcp.md does not state MAX_OPEN_SESSIONS=%d (%r)" + % (MAX_OPEN_SESSIONS, word) + ) + + def test_the_scan_result_cap_matches_the_default(self): + from PyMemoryEditor.mcp.config import DEFAULT_MAX_SCAN_RESULTS + + doc = self._mcp_doc() + # The docs group digits, as prose should: "100 000", not "100000". + grouped = "{:,}".format(DEFAULT_MAX_SCAN_RESULTS).replace(",", " ") + + assert "Default %s." % grouped in doc, ( + "docs/mcp.md does not state DEFAULT_MAX_SCAN_RESULTS=%s" + % grouped + ) + + class TestServerConfigValidatesItsBounds: """The three numeric bounds were checked at the CLI only. diff --git a/tests/mcp_server/test_server.py b/tests/mcp_server/test_server.py index 88cd631..58e56e8 100644 --- a/tests/mcp_server/test_server.py +++ b/tests/mcp_server/test_server.py @@ -437,6 +437,78 @@ def test_refusal_tells_the_model_not_to_go_looking_elsewhere(self): assert "did not approve" in text assert "Do not retry" in text + def test_a_full_server_refuses_before_spending_an_approval(self, monkeypatch): + """The cap lives in `SessionStore.open`, which runs after the handle is + already open — so the user used to be asked, approve, and only then be + told the server was full.""" + from PyMemoryEditor.mcp import session as session_module + from PyMemoryEditor.mcp.session import MAX_OPEN_SESSIONS + + # Fakes carry arbitrary pids, which the store's reaper reads as dead. + monkeypatch.setattr(session_module, "pid_exists", lambda pid: True) + + server, toolset, _process = self._build(ServerConfig()) + for _ in range(MAX_OPEN_SESSIONS): + toolset.store.open(FakeProcess(), 4242, "faketarget") + + answerer = Answerer("accept") + + async def run(): + async with connected(server, answerer) as session: + return await session.call_tool("open_process", {"pid": 4242}) + + result = asyncio.run(run()) + + assert result.is_error is True + assert answerer.asked is False, "the user was prompted for an attach that could not happen" + assert "close_process" in result.content[0].text + + def test_the_refusal_names_the_target_it_declined(self, monkeypatch): + """A model reads this and must not conclude the pid was the problem.""" + from PyMemoryEditor.mcp import session as session_module + from PyMemoryEditor.mcp.session import MAX_OPEN_SESSIONS + + # Fakes carry arbitrary pids, which the store's reaper reads as dead. + monkeypatch.setattr(session_module, "pid_exists", lambda pid: True) + + server, toolset, _process = self._build(ServerConfig()) + for _ in range(MAX_OPEN_SESSIONS): + toolset.store.open(FakeProcess(), 4242, "faketarget") + + async def run(): + async with connected(server, Answerer("accept")) as session: + return await session.call_tool("open_process", {"pid": 4242}) + + text = asyncio.run(run()).content[0].text + assert "faketarget" in text and "4242" in text + assert str(MAX_OPEN_SESSIONS) in text + # The early path renders the store's message, so it is as actionable + # as the late one: every open session, with its scan count. + assert "proc-1" in text and "scan" in text + + def test_an_unnamed_target_is_not_refused_as_an_empty_string(self, monkeypatch): + """`_name_for_pid` returns "" for a pid missing from the listing. + + The elicitation prompt below this already guards with `or "unknown"`; + the refusal did not, so it rendered `attach to "" (pid 4242)`. + """ + from PyMemoryEditor.mcp import session as session_module + from PyMemoryEditor.mcp.session import MAX_OPEN_SESSIONS + + monkeypatch.setattr(session_module, "pid_exists", lambda pid: True) + server, toolset, _process = self._build(ServerConfig()) + monkeypatch.setattr(toolset, "_name_for_pid", lambda pid: "") + for _ in range(MAX_OPEN_SESSIONS): + toolset.store.open(FakeProcess(), 4242, "faketarget") + + async def run(): + async with connected(server, Answerer("accept")) as session: + return await session.call_tool("open_process", {"pid": 4242}) + + text = asyncio.run(run()).content[0].text + assert 'attach to ""' not in text + assert "unknown" in text + def test_cancelling_the_dialog_is_not_approval(self): result, toolset = self._open(ServerConfig(), Answerer("cancel")) assert result.is_error is True diff --git a/tests/mcp_server/test_session.py b/tests/mcp_server/test_session.py index 42c72c8..89117db 100644 --- a/tests/mcp_server/test_session.py +++ b/tests/mcp_server/test_session.py @@ -6,8 +6,11 @@ import pytest +from PyMemoryEditor.mcp import session as session_module from PyMemoryEditor.mcp.session import ( + MAX_OPEN_SESSIONS, MAX_SCANS_PER_SESSION, + Session, SessionError, SessionStore, batch_regions, @@ -201,6 +204,328 @@ def test_snapshot_is_address_sorted(self, store): ) +@pytest.fixture +def all_alive(monkeypatch): + """Pin every session's target as live. + + The store reaps sessions whose pid is gone, and a fake's pid is an + arbitrary number — so without this these tests depend on the host's + process table, and the cap either fires or does not by luck. (Several + passed only because `pid_exists(1)` is true on a Unix host.) + """ + monkeypatch.setattr(session_module, "pid_exists", lambda pid: True) + + +class TestOpenSessionsAreCapped: + """The store used to accept any number of open processes, each holding an + OS handle that only `close_process` releases.""" + + def test_the_cap_is_enforced(self, store, all_alive): + for _ in range(MAX_OPEN_SESSIONS): + store.open(FakeProcess(), 1, "a") + + with pytest.raises(SessionError) as error: + store.open(FakeProcess(), 1, "a") + + assert str(MAX_OPEN_SESSIONS) in str(error.value) + + def test_the_refusal_says_how_to_recover(self, store, all_alive): + """The model cannot see the store, so the message has to name an id to + close and list what is open.""" + for index in range(MAX_OPEN_SESSIONS): + store.open(FakeProcess(pid=100 + index), 100 + index, "target%d" % index) + + with pytest.raises(SessionError) as error: + store.open(FakeProcess(), 999, "another") + + message = str(error.value) + assert "close_process" in message + assert "proc-1" in message # every open session is listed + assert "target0" in message + assert "pid 100" in message + # No "close the oldest": the oldest is usually the target the whole + # session has been refining, and closing it discards its scans too. + assert "oldest" not in message + # The scan count is what makes the list actionable. + assert "scan" in message + + def test_closing_one_frees_a_slot(self, store, all_alive): + ids = [store.open(FakeProcess(), 1, "a").session_id + for _ in range(MAX_OPEN_SESSIONS)] + + store.close(ids[0]) + reopened = store.open(FakeProcess(), 2, "b") + + assert reopened.session_id not in ids + assert len(store.sessions) == MAX_OPEN_SESSIONS + + def test_ids_keep_climbing_after_a_close(self, store, all_alive): + """Reusing an id would let a model holding a stale one address a + different process.""" + first = store.open(FakeProcess(), 1, "a").session_id + store.close(first) + second = store.open(FakeProcess(), 2, "b").session_id + + assert first != second + + +class TestTheCapHoldsUnderConcurrency: + """Tool calls run on a thread pool, so `open` races with itself. + + A refactor split the capacity check and the insert across two acquisitions + of the store's lock: every thread saw room, then every thread inserted. + The window is a few bytecodes wide, so it does not reproduce on its own — + hence the injected pause, which makes a structural defect observable + instead of waiting for luck. + """ + + def test_the_cap_is_not_exceeded_when_opens_race(self, store, all_alive): + import threading + + start = threading.Barrier(MAX_OPEN_SESSIONS * 2) + + def attach(): + start.wait() + try: + store.open(FakeProcess(), 1, "a") + except SessionError: + pass + + threads = [ + threading.Thread(target=attach) + for _ in range(MAX_OPEN_SESSIONS * 2) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert len(store.sessions) == MAX_OPEN_SESSIONS + + def test_capacity_is_judged_while_the_lock_is_held(self, store, all_alive): + """The property the timing test could not actually probe. + + A previous version of this test patched `capacity_refusal` and slept + in it to widen the race — but `open` calls `_refusal_locked`, not + `capacity_refusal`, so the sleep never fired and the test only + duplicated the one above. It passed its mutation check because that + mutant happened to reintroduce the `capacity_refusal` call. + + Asserted structurally instead: `_refusal_locked` refuses to judge + capacity unless the lock is held, so splitting the check from the + insert fails deterministically rather than when the scheduler + cooperates. + """ + for _ in range(MAX_OPEN_SESSIONS): + store.open(FakeProcess(), 1, "a") + + # Held: this is how `open` calls it, and it answers. + with store._lock: + assert store._refusal_locked(0) is not None + + # Not held: it refuses to answer at all. + with pytest.raises(AssertionError): + store._refusal_locked(0) + + def test_the_refusal_is_raised_inside_the_critical_section( + self, store, all_alive, monkeypatch + ): + """Deciding under the lock and raising outside it leaves a window. + + Another thread can close a session in that gap, so the caller is turned + away with a message that is already false — the store has room by the + time it reads it. Probed by recording whether the lock was held at the + moment the error was constructed. + """ + for _ in range(MAX_OPEN_SESSIONS): + store.open(FakeProcess(), 1, "a") + + held = [] + original = session_module.SessionError + + class Probe(original): # type: ignore[misc, valid-type] + def __init__(self, *args): + held.append(store._lock.locked()) + super().__init__(*args) + + monkeypatch.setattr(session_module, "SessionError", Probe) + + with pytest.raises(original): + store.open(FakeProcess(), 2, "b") + + assert held == [True], ( + "the refusal was built after the lock was released, so the store " + "may already have had room" + ) + + def test_open_judges_and_inserts_under_one_acquisition(self, all_alive): + """The assertion above is not enough on its own. + + It only proves capacity is judged under *an* acquisition, not that the + same one covers the insert — a version that judged under its own lock, + released, and then took the lock again to insert passed it, and that is + exactly the regression this guards. Counting acquisitions separates + them: one for the correct implementation, two for the split. + """ + import threading + + class CountingLock: + def __init__(self): + self._real = threading.Lock() + self.acquisitions = 0 + + def __enter__(self): + self.acquisitions += 1 + return self._real.__enter__() + + def __exit__(self, *exc): + return self._real.__exit__(*exc) + + def locked(self): + return self._real.locked() + + store = SessionStore() + lock = CountingLock() + store._lock = lock + + store.open(FakeProcess(), 1, "a") + + assert lock.acquisitions == 1, ( + "capacity and the insert must share one acquisition; %d means the " + "check released the lock before inserting" % lock.acquisitions + ) + + def test_a_dead_session_is_reaped_even_if_the_store_fills_meanwhile( + self, store, monkeypatch + ): + """`open` used to decide whether to reap from a check taken under an + earlier acquisition of the lock. If the store filled in between, it + refused without anyone having looked for a dead session. + + Simulated by filling the store from inside the first check, which is + the same ordering a concurrent open produces. + """ + alive = {4242} + monkeypatch.setattr(session_module, "pid_exists", lambda pid: pid in alive) + + store.open(FakeProcess(pid=4242), 4242, "live") + + original = type(store)._refusal_locked + calls = [] + + def fill_once(self, pid): + if not calls: + calls.append(pid) + for index in range(MAX_OPEN_SESSIONS - 1): + session_id = "proc-%d" % next(self._session_ids) + self._sessions[session_id] = Session( + session_id=session_id, process=FakeProcess(pid=700 + index), + pid=700 + index, name="gone%d" % index, + ) + return original(self, pid) + + monkeypatch.setattr(type(store), "_refusal_locked", fill_once) + + # Full at the moment of the check, but seven of the eight are dead. + session = store.open(FakeProcess(pid=4243), 4243, "second") + + assert session.pid == 4243 + + +class TestDeadSessionsDoNotHoldSlots: + """The lockout the cap created, and the reason it evicts here but nowhere + else. + + Nothing in this server notices a target exiting: `_rows_for` swallows the + read error, `process_info` keeps answering from cached state, and the + session stays in the store. With a cap that refuses rather than evicts, + eight exited processes held every slot for the life of the server and the + model had no way to tell which sessions were the dead ones. + + Evicting a *dead* session is the exception that proves the rule about not + evicting: its handle is already useless, so nothing is taken away. + """ + + def test_a_dead_session_is_reaped_to_make_room(self, store, monkeypatch): + for index in range(MAX_OPEN_SESSIONS): + store.open(FakeProcess(pid=500 + index), 500 + index, "gone%d" % index) + + monkeypatch.setattr(session_module, "pid_exists", lambda pid: pid == 4242) + + # Would have raised before the reaper existed. + session = store.open(FakeProcess(pid=4242), 4242, "live") + + assert session.pid == 4242 + assert [held.pid for held in store.sessions] == [4242], ( + "the eight dead sessions should be gone, not merely joined" + ) + + def test_a_live_session_is_never_reaped(self, store, monkeypatch): + store.open(FakeProcess(pid=4242), 4242, "live") + for index in range(MAX_OPEN_SESSIONS - 1): + store.open(FakeProcess(pid=600 + index), 600 + index, "gone%d" % index) + + monkeypatch.setattr(session_module, "pid_exists", lambda pid: pid == 4242) + store.open(FakeProcess(pid=4243), 4243, "second") + + pids = sorted(s.pid for s in store.sessions) + assert 4242 in pids, "a live session was evicted" + + def test_reaping_closes_the_handle(self, store, monkeypatch): + process = FakeProcess(pid=700) + store.open(process, 700, "gone") + monkeypatch.setattr(session_module, "pid_exists", lambda pid: False) + + store._reap_dead() + + assert process.closed is True, "the handle was dropped without closing" + assert store.sessions == () + + @pytest.mark.parametrize("pid", [2 ** 31, 2 ** 64, -3, 0]) + def test_a_pid_the_os_refuses_reads_as_dead_not_as_a_crash(self, store, pid): + """The reaper runs inside `open`, so anything it raises escapes as + something other than SessionError and the model sees only "Error + executing tool". + + `pid_exists(2**31)` raises OverflowError — the value does not fit the + platform's pid type — which the argument fuzzer found by generating + exactly that. Intermittently, because it also needed the store to be + at capacity at that moment. + """ + from PyMemoryEditor.mcp.session import looks_alive + + assert looks_alive(pid) is False + + for index in range(MAX_OPEN_SESSIONS): + store.open(FakeProcess(pid=pid), pid, "weird%d" % index) + + # Must not raise anything but SessionError, and here it makes room. + session = store.open(FakeProcess(pid=pid), pid, "another") + assert session.pid == pid + + def test_a_full_server_of_live_targets_still_refuses(self, store, all_alive): + """The reaper must not become a back door around the cap.""" + for index in range(MAX_OPEN_SESSIONS): + store.open(FakeProcess(pid=800 + index), 800 + index, "live%d" % index) + + with pytest.raises(SessionError) as error: + store.open(FakeProcess(pid=900), 900, "another") + + assert "all of them are live" in str(error.value) + + def test_the_refusal_points_at_an_already_open_pid(self, store, all_alive): + """A model that lost count re-attaches instead of reusing the id.""" + store.open(FakeProcess(pid=4242), 4242, "faketarget") + for index in range(MAX_OPEN_SESSIONS - 1): + store.open(FakeProcess(pid=810 + index), 810 + index, "other%d" % index) + + with pytest.raises(SessionError) as error: + store.open(FakeProcess(pid=4242), 4242, "faketarget") + + message = str(error.value) + assert "already open as proc-1" in message + + class TestBatchRegions: def _regions(self, sizes): address = 0x1000 @@ -221,28 +546,18 @@ def test_batches_respect_the_byte_budget(self): assert len(batches) == 4 # 3 x 300 bytes, then the 100-byte remainder def test_a_region_larger_than_the_budget_gets_its_own_batch(self): - # It cannot be split without splitting a value across the seam, so the - # budget is a target rather than a guarantee -- but the batch that - # busts it should not be carrying anything else. - # - # This assertion used to be `[2, 1]`, i.e. the oversized region shared - # a batch with the small one before it, under this same name. The name - # was right and the assertion pinned the opposite, so the test was - # documenting the bug it looked like it was guarding against. + # A region can't be split without splitting a value across the seam, + # so the budget is a target -- but the batch that busts it should not + # carry anything else. This asserted `[2, 1]` under the same name, + # i.e. it pinned the opposite of what the name claims. batches = batch_regions(self._regions([10, 5000, 10]), 100) assert [[region.size for region in batch] for batch in batches] == [ [10], [5000], [10] ] def test_a_run_of_small_regions_does_not_ride_along_with_a_large_one(self): - """The case that made the deadline check pointless. - - `batch_regions` exists so `_run_batched_scan` gets a chance to look at - the clock between batches. Without the flush, every small region before - an oversized one joined it: `[10, 10, 10, 5000]` against a 100-byte - budget came back as a *single* batch of four, so the deadline was - checked once for the whole scan -- the one thing the batching is for. - """ + """Without the flush, `[10, 10, 10, 5000]` on a 100-byte budget was one + batch of four, so the deadline was checked once for the whole scan.""" batches = batch_regions(self._regions([10, 10, 10, 5000]), 100) assert len(batches) == 2 diff --git a/tests/mcp_server/test_toolset.py b/tests/mcp_server/test_toolset.py index e74d652..91c1ff1 100644 --- a/tests/mcp_server/test_toolset.py +++ b/tests/mcp_server/test_toolset.py @@ -1210,6 +1210,68 @@ def test_max_offset_is_still_clamped_at_the_top(self, toolset, session): # --- the one address no parser ever sees --------------------------- # + def test_server_info_survives_a_session_with_an_impossible_pid( + self, make_toolset + ): + """`server_info` is the tool the model is told to call first. + + It rendered `alive` with a bare `pid_exists`, which raises + OverflowError for a pid outside the platform's range — so one odd + session took the whole result down, hiding the limits and the policy + as well as the session list. The store had already been guarded for + the same call; this file had not. + """ + # `os.getpid()`, not pid 1: Windows has no pid 1 (System Idle is 0, + # System is 4), so `pid_exists(1)` is False there and this test failed + # on that runner alone. The test's own process is alive by definition + # everywhere. + import os + + alive_pid = os.getpid() + + toolset = make_toolset(config()) + toolset.store.open(FakeProcess(pid=2 ** 31), 2 ** 31, "impossible") + toolset.store.open(FakeProcess(pid=alive_pid), alive_pid, "ordinary") + + sessions = toolset.server_info()["open_sessions"] + + by_pid = {entry["pid"]: entry["alive"] for entry in sessions} + assert by_pid[2 ** 31] is False + assert by_pid[alive_pid] is True + + def test_a_refused_attach_does_not_leak_the_handle( + self, make_toolset, monkeypatch + ): + """`attach` opens the handle and *then* registers the session, so a + refusal must close it or the cap leaks what it exists to bound.""" + from PyMemoryEditor.mcp import session as session_module + from PyMemoryEditor.mcp.session import MAX_OPEN_SESSIONS + + # The store reaps sessions whose pid is gone, and a fake's pid is an + # arbitrary number, so the cap would never fire here otherwise. + monkeypatch.setattr(session_module, "pid_exists", lambda pid: True) + + opened = [] + + def fake_open(**_kwargs): + process = FakeProcess() + opened.append(process) + return process + + toolset = make_toolset(config()) + monkeypatch.setattr(toolset, "_open_process", fake_open) + + for _ in range(MAX_OPEN_SESSIONS): + toolset.open_process(pid=4242) + + with pytest.raises(SessionError): + toolset.open_process(pid=4242) + + # The last one is the refused attach's, and only it should be closed. + assert len(opened) == MAX_OPEN_SESSIONS + 1 + assert opened[-1].closed is True + assert all(process.closed is False for process in opened[:-1]) + def test_a_chain_resolving_past_64_bits_is_refused_not_returned( self, make_toolset, fake_process ):