Skip to content

Commit 59c96b9

Browse files
authored
[FIX]: Fix async context manager protocol typing (#119)
## Description `ty` (>=0.59.0, as seen in #118) rejects `Session` and `InjectionHandle` when they are passed to `AsyncExitStack.enter_async_context`. Both protocols inherit their context manager methods from `AbstractAsyncContextManager`, but also redeclare `__aenter__` using `Self`. That leaves `ty` trying to reconcile two recursive versions of the same signature. This removes the duplicate `__aenter__` and `__aexit__` declarations and relies on the inherited signatures instead. `Session` now specifies `None` as its exit return type, preserving the previous cleanup contract. The XPIA call sites do not need to change. Validation: - `uv run ty check` - `uv run pytest tests/unit/core/test_adapter.py tests/unit/core/test_injection.py tests/unit/attacks/test_xpia.py` - `uv run pre-commit run --all-files` ## Breaking changes None. ## Checklist - [x] `pre-commit run --all-files` passes - [ ] Tests added or updated for changes (existing protocol and XPIA tests cover the affected behavior) - [x] Documentation updated
1 parent 1a67e72 commit 59c96b9

2 files changed

Lines changed: 8 additions & 39 deletions

File tree

rampart/core/adapter.py

Lines changed: 4 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -10,22 +10,20 @@
1010
from __future__ import annotations
1111

1212
from contextlib import AbstractAsyncContextManager
13-
from typing import TYPE_CHECKING, Protocol, Self, runtime_checkable
13+
from typing import TYPE_CHECKING, Protocol, runtime_checkable
1414

1515
if TYPE_CHECKING:
16-
import types
17-
1816
from rampart.core.manifest import AppManifest
1917
from rampart.core.types import ObservabilityLevel, Request, Response
2018

2119

2220
@runtime_checkable
23-
class Session(AbstractAsyncContextManager["Session"], Protocol):
21+
class Session(AbstractAsyncContextManager["Session", None], Protocol):
2422
"""A bounded unit of interaction with the agent.
2523
2624
Sessions are async context managers. Entering returns the session
27-
ready for use; exiting guarantees cleanup of any resources the
28-
adapter holds (API clients, browser contexts, temporary state).
25+
ready for use; exiting guarantees idempotent cleanup of any resources
26+
the adapter holds (API clients, browser contexts, temporary state).
2927
3028
Fresh state = fresh session. Create a new one via the adapter.
3129
"""
@@ -46,20 +44,6 @@ async def send_async(self, request: Request) -> Response:
4644
"""
4745
...
4846

49-
async def __aenter__(self) -> Self:
50-
"""Enter the session context. Returns self."""
51-
...
52-
53-
async def __aexit__(
54-
self,
55-
exc_type: type[BaseException] | None,
56-
exc_value: BaseException | None,
57-
traceback: types.TracebackType | None,
58-
/,
59-
) -> None:
60-
"""Clean up session resources. Must be idempotent."""
61-
...
62-
6347

6448
@runtime_checkable
6549
class AgentAdapter(Protocol):

rampart/core/injection.py

Lines changed: 4 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,9 @@
1414

1515
import asyncio
1616
from contextlib import AbstractAsyncContextManager
17-
from typing import TYPE_CHECKING, Protocol, Self, runtime_checkable
17+
from typing import TYPE_CHECKING, Protocol, runtime_checkable
1818

1919
if TYPE_CHECKING:
20-
import types
21-
2220
from rampart.core.types import Payload
2321

2422

@@ -27,8 +25,9 @@ class InjectionHandle(AbstractAsyncContextManager["InjectionHandle", None], Prot
2725
"""A prepared injection, ready to activate as an async context manager.
2826
2927
Returned by Surface.inject(). Entering activates the injection
30-
(writes the payload to the data source); exiting removes it
31-
(guaranteed cleanup even on exceptions).
28+
(writes the payload to the data source); exiting removes it.
29+
Cleanup is guaranteed even on exceptions, must be idempotent, and
30+
must not raise.
3231
3332
Execution strategies depend only on this protocol — never on
3433
Surface or its concrete implementations.
@@ -52,20 +51,6 @@ async def wait_until_ready(self) -> None:
5251
"""
5352
...
5453

55-
async def __aenter__(self) -> Self:
56-
"""Activate the injection (write payload to data source)."""
57-
...
58-
59-
async def __aexit__(
60-
self,
61-
exc_type: type[BaseException] | None,
62-
exc_value: BaseException | None,
63-
traceback: types.TracebackType | None,
64-
/,
65-
) -> None:
66-
"""Remove the injection. Must be idempotent. Must not raise."""
67-
...
68-
6954

7055
async def sleep_until_ready(delay: float) -> None:
7156
"""Sleep for `delay` seconds. Default readiness strategy for simple surfaces.

0 commit comments

Comments
 (0)