Stop new Notion writes after cancellation and bound response reads

davellm_notion.py

@@ -22,10 +22,11 @@
   a fresh read could not settle it. An append with this outcome is never repeated by
   the adapter, and the run ledger refuses the same append again for the rest of the run.
 
-The write and its verification run in a task shielded from cancellation, so a run
-deadline can stop waiting for it but cannot cut it off between the request and the
-check; the late outcome is still recorded in the run ledger. Writes to one page are
-serialized within the process, so a match found while checking an uncertain append
+Write verification runs in a task shielded from cancellation. A cancelled caller
+prevents any new write request, including a retry after rate limiting; a request
+already dispatched can still be checked, with the late append outcome recorded in
+the run ledger. Each HTTP request has an absolute deadline covering its entire body.
+Writes to one page are serialized within the process, so a match found while checking an uncertain append
 cannot be another DaveLLM run's identical write.
 
 Transport. Requests go only to ``https://api.notion.com/v1`` with a pinned
@@ -46,6 +47,7 @@
 import threading
 import time
 from collections import OrderedDict
+from contextvars import ContextVar
 from dataclasses import dataclass, field
 from typing import Any, Coroutine, Mapping, Optional, Sequence
 from urllib.parse import urlparse
@@ -337,6 +339,8 @@ def busy(self, page_id: str) -> bool:
 
 # Shielded write tasks outlive a cancelled caller; hold them so they are not collected.
 _WRITES_IN_FLIGHT: set[asyncio.Task[Any]] = set()
+# Shared with the shielded task: cancellation prevents new writes but allows readback.
+_WRITE_CANCELLED: ContextVar[Optional[asyncio.Event]] = ContextVar("notion_write_cancelled", default=None)
 
 
 def _copy_json(value: Any) -> Any:
@@ -581,9 +585,17 @@ async def call(self, method: str, path: str, *, params: Optional[dict[str, Any]]
             remaining = self.deadline - time.monotonic()
             if remaining < NOTION_MIN_REQUEST_SECONDS or self.requests >= self.max_requests:
                 raise _Budget()
+            cancelled = _WRITE_CANCELLED.get()
+            if method != "GET" and cancelled is not None and cancelled.is_set():
+                raise _Budget()
             self.requests += 1
-            status, headers, payload = await self._send(method, path, params, body,
-                                                        min(NOTION_REQUEST_TIMEOUT_SECONDS, remaining))
+            timeout = min(NOTION_REQUEST_TIMEOUT_SECONDS, remaining)
+            try:
+                status, headers, payload = await asyncio.wait_for(
+                    self._send(method, path, params, body, timeout), timeout=timeout)
+            except asyncio.TimeoutError:
+                # A request may have reached Notion; do not label an expired write unsent.
+                raise _Uncertain() from None
             if 200 <= status < 300:
                 if not isinstance(payload, dict):
                     raise _Uncertain()
@@ -601,6 +613,9 @@ async def call(self, method: str, path: str, *, params: Optional[dict[str, Any]]
 
     async def _send(self, method: str, path: str, params: Optional[dict[str, Any]],
                     body: Optional[dict[str, Any]], timeout: float) -> tuple[int, httpx.Headers, Any]:
+        cancelled = _WRITE_CANCELLED.get()
+        if method != "GET" and cancelled is not None and cancelled.is_set():
+            raise _Budget()
         try:
             async with self._client.stream(method, path, params=params, json=body,
                                            timeout=httpx.Timeout(timeout)) as response:
@@ -780,16 +795,24 @@ async def _shielded(guard: PageWriteGuard, page_id: str, work: Coroutine[Any, An
         work.close()
         raise NotionToolError(PAGE_BUSY)
 
+    cancelled = asyncio.Event()
+
     async def guarded() -> _Outcome:
+        token = _WRITE_CANCELLED.set(cancelled)
         try:
             return await work
         finally:
+            _WRITE_CANCELLED.reset(token)
             guard.release(page_id)
 
     task = asyncio.ensure_future(guarded())
     _WRITES_IN_FLIGHT.add(task)
     task.add_done_callback(_WRITES_IN_FLIGHT.discard)
-    return await asyncio.shield(task)
+    try:
+        return await asyncio.shield(task)
+    except asyncio.CancelledError:
+        cancelled.set()
+        raise
 
 
 def _append_blocks_argument(value: Any) -> list[tuple[str, str, Optional[bool]]]:

tests/test_notion_independent_review.py

@@ -0,0 +1,145 @@
+"""Independent regression probes; all network traffic is simulated."""
+import asyncio
+
+import httpx
+import pytest
+import respx
+import davellm_notion as notion
+from test_notion_tools import World, BLOCKS
+
+
+def test_cancel_before_dispatch_must_prevent_write(monkeypatch):
+    world = World()
+    original = notion.NotionApi.call
+
+    async def scenario():
+        entered = asyncio.Event()
+        release = asyncio.Event()
+        async def paused(self, method, path, **kwargs):
+            if method == 'GET' and path.endswith('/children') and not entered.is_set():
+                entered.set()
+                await release.wait()
+            return await original(self, method, path, **kwargs)
+        monkeypatch.setattr(notion.NotionApi, 'call', paused)
+        with respx.mock(assert_all_mocked=True, assert_all_called=False) as mock:
+            world.fake.mount(mock)
+            caller = asyncio.create_task(notion.append_blocks(
+                {'page': 'adapter-test', 'blocks': BLOCKS}, settings=world.settings,
+                ledger=world.ledger, guard=world.guard))
+            await asyncio.wait_for(entered.wait(), 1)
+            caller.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await caller
+            release.set()
+            pending = list(notion._WRITES_IN_FLIGHT)
+            if pending:
+                await asyncio.wait_for(asyncio.gather(*pending), 1)
+            assert len(world.fake.writes()) == 0, 'PATCH dispatched after caller cancellation, during preflight GET'
+    asyncio.run(scenario())
+
+
+def test_request_budget_is_absolute_during_slow_response(monkeypatch):
+    monkeypatch.setattr(notion, 'NOTION_MIN_REQUEST_SECONDS', 0)
+    class SlowBody(httpx.AsyncByteStream):
+        async def __aiter__(self):
+            for chunk in (b'{', b'"ok"', b':true', b'}'):
+                await asyncio.sleep(.04)
+                yield chunk
+    async def scenario():
+        with respx.mock(assert_all_mocked=True, assert_all_called=False) as mock:
+            mock.get('https://api.notion.com/v1/users/me').mock(
+                return_value=httpx.Response(200, stream=SlowBody()))
+            async with notion.NotionApi('fake-token', budget_seconds=.06, max_requests=1) as api:
+                with pytest.raises(notion._Uncertain):
+                    await api.call('GET', '/users/me')
+    asyncio.run(scenario())
+
+
+def test_cancel_after_dispatch_allows_readback_without_second_write(monkeypatch):
+    world = World()
+    original = notion.NotionApi.call
+    async def scenario():
+        sent = asyncio.Event()
+        release = asyncio.Event()
+        async def paused(self, method, path, **kwargs):
+            result = await original(self, method, path, **kwargs)
+            if method == 'PATCH':
+                sent.set()
+                await release.wait()
+            return result
+        monkeypatch.setattr(notion.NotionApi, 'call', paused)
+        with respx.mock(assert_all_mocked=True, assert_all_called=False) as mock:
+            world.fake.mount(mock)
+            caller = asyncio.create_task(notion.append_blocks(
+                {'page': 'adapter-test', 'blocks': BLOCKS}, settings=world.settings,
+                ledger=world.ledger, guard=world.guard))
+            await asyncio.wait_for(sent.wait(), 1)
+            caller.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await caller
+            release.set()
+            await asyncio.wait_for(asyncio.gather(*list(notion._WRITES_IN_FLIGHT)), 1)
+            assert len(world.fake.writes()) == 1
+            assert set(world.ledger._appends.values()) == {'verified'}
+            assert not world.guard.busy(world.page)
+    asyncio.run(scenario())
+
+
+def test_cancel_before_update_dispatch_must_prevent_write(monkeypatch):
+    from fake_notion import rt
+    world = World()
+    block_id = world.fake.add_block(world.page, 'paragraph', [rt('before')])
+    world.read()
+    original = notion.NotionApi.call
+    async def scenario():
+        entered = asyncio.Event()
+        release = asyncio.Event()
+        async def paused(self, method, path, **kwargs):
+            if method == 'GET' and path == f'/blocks/{block_id}' and not entered.is_set():
+                entered.set()
+                await release.wait()
+            return await original(self, method, path, **kwargs)
+        monkeypatch.setattr(notion.NotionApi, 'call', paused)
+        with respx.mock(assert_all_mocked=True, assert_all_called=False) as mock:
+            world.fake.mount(mock)
+            caller = asyncio.create_task(notion.update_block(
+                {'page': 'adapter-test', 'block': world.ref(block_id),
+                 'old_text': 'before', 'new_text': 'after'},
+                settings=world.settings, ledger=world.ledger, guard=world.guard))
+            await asyncio.wait_for(entered.wait(), 1)
+            caller.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await caller
+            release.set()
+            await asyncio.wait_for(asyncio.gather(*list(notion._WRITES_IN_FLIGHT)), 1)
+            assert len(world.fake.writes()) == 0
+            assert world.fake.text_of(block_id) == 'before'
+    asyncio.run(scenario())
+
+
+def test_cancel_during_rate_limit_wait_prevents_write_retry(monkeypatch):
+    world = World()
+    world.fake.fail("PATCH", r"/blocks/.*/children", "rate_limit", retry_after="0.1")
+    original = notion.NotionApi._send
+    async def scenario():
+        rate_limited = asyncio.Event()
+        async def observe(self, method, path, params, body, timeout):
+            result = await original(self, method, path, params, body, timeout)
+            if result[0] == 429:
+                rate_limited.set()
+            return result
+        monkeypatch.setattr(notion.NotionApi, "_send", observe)
+        with respx.mock(assert_all_mocked=True, assert_all_called=False) as mock:
+            world.fake.mount(mock)
+            caller = asyncio.create_task(notion.append_blocks(
+                {"page": "adapter-test", "blocks": BLOCKS}, settings=world.settings,
+                ledger=world.ledger, guard=world.guard))
+            await asyncio.wait_for(rate_limited.wait(), 1)
+            caller.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await caller
+            await asyncio.wait_for(asyncio.gather(*list(notion._WRITES_IN_FLIGHT)), 1)
+            assert len(world.fake.writes()) == 1
+            assert world.fake.page_texts(world.page) == []
+            assert set(world.ledger._appends.values()) == {"failed"}
+    asyncio.run(scenario())