diff --git a/engineering/pid/spec.py b/engineering/pid/spec.py index db45ff9..da335a6 100644 --- a/engineering/pid/spec.py +++ b/engineering/pid/spec.py @@ -6,6 +6,8 @@ import copy import math from typing import TYPE_CHECKING +import anyio + from .critique import PID_FOCUSES, issues_for from .drawlines import draw_line from .graph import build_graph @@ -443,7 +445,18 @@ async def run_spec(backend: AutoCADBackend, spec: dict, dry_run: bool = False) - # a ``BaseException``) leave the placed symbols in the drawing inside # an open transaction — on COM with the undo mark still open, so every # later ``pid_from_spec`` was refused until a manual rollback. - await backend.transaction_rollback() + # + # The rollback is shielded because the MCP SDK cancels a request by + # cancelling an anyio scope, and anyio cancellation is level-triggered: + # ``task.cancel()`` is re-issued on every loop iteration while the task + # is still inside the cancelled scope. A bare ``await`` here would be + # cancelled at its first suspension, before the rollback body ran — + # measured through the real transport: five symbols left in an open + # transaction, and on COM the undo mark left open while the backend's + # flag already said "no transaction". A native ``Task.cancel()`` is + # delivered once and clears, so it was never the case that mattered. + with anyio.CancelScope(shield=True): + await backend.transaction_rollback() raise await backend.transaction_commit() graph = await build_graph(backend) diff --git a/tests/test_pid_from_spec.py b/tests/test_pid_from_spec.py index 934dea4..32247df 100644 --- a/tests/test_pid_from_spec.py +++ b/tests/test_pid_from_spec.py @@ -110,20 +110,180 @@ async def test_non_finite_numbers_are_refused_by_path_before_any_write(backend, assert (await backend.system_status())["transaction_depth"] == 0 -async def test_a_cancellation_mid_run_rolls_the_sheet_back(backend, monkeypatch): - """``CancelledError`` is a ``BaseException``; ``except Exception`` let the - placed symbols stay inside an open transaction.""" +async def _wait_until(predicate, *, timeout: float = 5.0) -> None: + """Poll an async predicate; a rollback that never lands fails here, not by hanging.""" import asyncio + import time + + deadline = time.monotonic() + timeout + while not await predicate(): + if time.monotonic() > deadline: + raise AssertionError(f"condition not met within {timeout:.1f}s") + await asyncio.sleep(0.02) + + +@pytest.mark.parametrize("how", ["scope", "native"]) +async def test_a_cancellation_mid_run_rolls_the_sheet_back(backend, monkeypatch, how): + """A *real* cancellation, not a ``CancelledError`` raised from a mock. + + The MCP SDK cancels a request through an anyio cancel scope, which is + level-triggered: ``task.cancel()`` is re-issued on every loop iteration + while the task is inside the cancelled scope, so an unshielded rollback + was cancelled at its first ``await`` before its body ran (five symbols + left in an open transaction). A native ``Task.cancel()`` is delivered once + and clears — raising from the mock only ever exercised that easy case. + """ + import asyncio + + import anyio from engineering.pid import spec as spec_module before = await backend.entity_count() + in_flight = asyncio.Event() - async def cancelled(*args, **kwargs): - raise asyncio.CancelledError() + async def slow_draw_line(*args, **kwargs): + in_flight.set() + await asyncio.sleep(30) + + monkeypatch.setattr(spec_module, "draw_line", slow_draw_line) + + if how == "scope": + async with anyio.create_task_group() as tg: + tg.start_soon(run_spec, backend, EXAMPLE_SPEC) + await in_flight.wait() + assert await backend.entity_count() > before, "symbols placed before the cancel" + tg.cancel_scope.cancel() + else: + task = asyncio.ensure_future(run_spec(backend, EXAMPLE_SPEC)) + await in_flight.wait() + assert await backend.entity_count() > before, "symbols placed before the cancel" + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task - monkeypatch.setattr(spec_module, "draw_line", cancelled) - with pytest.raises(asyncio.CancelledError): - await run_spec(backend, EXAMPLE_SPEC) assert await backend.entity_count() == before, "the placed symbols were rolled back" assert (await backend.system_status())["transaction_depth"] == 0, "no open transaction" + + +async def test_a_client_cancellation_through_the_transport_rolls_the_sheet_back( + client, monkeypatch +): + """The docstring's promise, measured end to end: ``notifications/cancelled`` + from a real ``fastmcp.Client`` mid-run leaves no half-drawn sheet and no + open transaction.""" + import asyncio + + from fastmcp.exceptions import McpError + + from engineering.pid import spec as spec_module + + async def slow_draw_line(*args, **kwargs): + await asyncio.sleep(30) + + monkeypatch.setattr(spec_module, "draw_line", slow_draw_line) + + async def depth() -> int: + return (await client.call_tool("system_status", {})).structured_content["transaction_depth"] + + async def entities() -> int: + return (await client.call_tool("analysis_entity_stats", {})).structured_content[ + "total_entities" + ] + + before = await entities() + # The session hands out ids sequentially, so the next one is the tool + # call's — provided nothing else is sent before the call goes out. Yield + # (no requests) until the counter has moved, then start polling. + request_id = client.session._request_id + call = asyncio.ensure_future(client.call_tool("pid_from_spec", {"spec": EXAMPLE_SPEC})) + + async def call_sent() -> bool: + return client.session._request_id > request_id + + await _wait_until(call_sent) + await _wait_until(lambda: _is_mid_run(depth, entities, before)) + await client.cancel(request_id) + with pytest.raises(McpError, match="cancelled"): + await call + # The SDK answers the cancel before the handler's rollback has landed. + await _wait_until(lambda: _is_clean(depth, entities, before)) + + +async def _is_mid_run(depth, entities, before) -> bool: + return await depth() == 1 and await entities() > before + + +async def _is_clean(depth, entities, before) -> bool: + return await depth() == 0 and await entities() == before + + +async def test_a_cancellation_with_a_com_call_in_flight_still_closes_the_undo_mark( + monkeypatch, +): + """COM shape: real ``ComBackend._run``, a real single-thread executor and the + real ``transaction_rollback``; only the ActiveX document is a recorder. + + With a COM call blocking the worker when the scope is cancelled, the + unshielded rollback's ``_run`` was cancelled while queued behind it — so + ``EndUndoMark`` / ``_UNDO B`` never reached AutoCAD — yet its ``finally`` + cleared ``_transaction_active``: ``system_status`` said "no transaction" + while the undo mark was open, and the next ``transaction_begin`` nested a + second ``StartUndoMark``. + """ + import asyncio + import time + import types + from concurrent.futures import ThreadPoolExecutor + + import anyio + + from backends import com_backend as module + from engineering.pid import spec as spec_module + + sent: list[str] = [] + doc = types.SimpleNamespace( + Name="Drawing1.dwg", + FullName="", + StartUndoMark=lambda: sent.append("StartUndoMark"), + EndUndoMark=lambda: sent.append("EndUndoMark"), + ) + app = types.SimpleNamespace(Documents=types.SimpleNamespace(Count=1), ActiveDocument=doc) + monkeypatch.setattr(module, "_acad_app", lambda: app) # _ensure_document_state + monkeypatch.setattr(module, "_acad_doc", lambda: doc) # begin / rollback + monkeypatch.setattr( + module.ComBackend, + "_safe_send_command", + staticmethod(lambda doc, cmd, deadline_s=8.0: sent.append(cmd) or []), + ) + backend = module.ComBackend() + backend._executor = ThreadPoolExecutor(max_workers=1) # no CoInitialize: nothing is COM here + try: + placed = 0 + + async def fake_place(backend, section, item): + nonlocal placed + placed += 1 + return {"handle": f"H{placed}"} + + in_flight = asyncio.Event() + + async def slow_draw_line(backend, *args, **kwargs): + in_flight.set() + await backend._run(time.sleep, 0.5) # the worker is inside a COM call + + monkeypatch.setattr(spec_module, "_place", fake_place) + monkeypatch.setattr(spec_module, "draw_line", slow_draw_line) + + async with anyio.create_task_group() as tg: + tg.start_soon(run_spec, backend, EXAMPLE_SPEC) + await in_flight.wait() + tg.cancel_scope.cancel() + + assert sent == ["StartUndoMark", "EndUndoMark", "_UNDO B"], sent + assert backend._transaction_active is False + # The mark really is closed, so the next transaction is a fresh one. + assert (await backend.transaction_begin())["ok"] is True + assert sent[-1] == "StartUndoMark" + finally: + backend._executor.shutdown(wait=True)