diff --git a/CHANGELOG.md b/CHANGELOG.md index dbff84a..4874a4d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -151,6 +151,20 @@ This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). no schema can be found the vectors report as `skip` with the flag to pass; `--only schema` exits non-zero, so a skip cannot quietly pass for a pass. +- **Scoring edits are rate-limited per cell** (`amp_server.ratelimit`). + RFC-AMP-001 §5 lists decay-score manipulation as a threat and asks + implementations to rate-limit scoring PATCHes per cell; nothing did. A caller + looping on `scoring` can hold a cell `active` past its relevance window or drive + a competing memory into archive. The budget is per cell, counts only PATCHes that + actually carry `scoring` (refusing ordinary content edits to cover an attack they + have nothing to do with would be an outage, not a mitigation), and answers + `429 RATE_LIMITED` with `Retry-After`. A refused attempt is not recorded, so + hammering cannot push back the caller's own deadline. Access is checked first, so + a caller who may not write the cell learns nothing about the remaining budget. + Default 5 per cell per hour, changed with `AMP_SCORING_PATCH_LIMIT` / + `AMP_SCORING_PATCH_WINDOW_SECONDS`, disabled with `0`, and advertised at + `GET /spec` as `scoring_patch_limit` (`null` when off). Counters are per process. + ### Changed - **Every endpoint returns one error shape.** `PATCH /memories/{id}` answered a conflict with `{"detail": ...}` while `DELETE` answered with diff --git a/docs/api-reference.md b/docs/api-reference.md index 3532a84..922898b 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -136,6 +136,7 @@ curl http://localhost:8765/amp/v1/spec "mcp_compatible": false, "storage_backends": ["chroma"], "api_keys_required": false, + "scoring_patch_limit": {"max_patches": 5, "window_seconds": 3600}, "embedding": {"provider": "chroma-default", "dimensions": 384}, "max_cell_size_bytes": 65536, "retention_days": 30, @@ -158,6 +159,7 @@ suite checks it against the server's own numbers rather than a fixed value: | Capability | What it commits the server to | |---|---| | `mcp_compatible` | Whether **this HTTP server** speaks MCP directly. It is `false`: the MCP integration ships as a separate stdio process (`amp-mcp`, see `examples/mcp-claude-desktop/`), not as an endpoint on this API. | +| `scoring_patch_limit` | How often one cell's `scoring` may be rewritten, and over what window (`null` when the limit is off). Enforced on `PATCH`: RFC-AMP-001 §5 names decay-score manipulation as a threat, because a caller looping on `scoring` can hold a cell `active` past its relevance window or push a competing memory into archive. A refusal is `429 RATE_LIMITED` with `Retry-After`. Section `PATCH /memories/{memory_id}` below covers what is and is not counted. | | `max_page_size` | The largest `limit` the listing endpoints accept (`100`). A larger value is refused with `422` rather than silently clamped, so a client never believes it received a complete page when it did not. | | `api_keys_required` | Whether this server requires `X-AMP-API-Key` (`AMP_API_KEYS_FILE` is set). Reported here so a client learns it needs a key before a call fails with `401`. | | `storage_backends` | The adapter actually wired in (`chroma` or `postgres`), selected with `AMP_STORAGE_BACKEND`; see [getting started](getting-started.md#6-choosing-a-storage-backend). | @@ -451,6 +453,34 @@ Fields that cannot be patched: `id`, `amp_version`, `identity`, `lifecycle.creat --- +### Scoring updates are rate-limited per cell + +RFC-AMP-001 §5 lists decay-score manipulation as a threat: a caller PATCHing +`scoring` in a loop can keep a cell `active` past its intended relevance window, +or force a competing memory into archive. Two things bound it. A scoring change +only takes effect on the next lifecycle pass, so the engine's cadence limits how +fast a manipulation lands; and the server budgets how often one cell's `scoring` +may be rewritten at all. + +- Only a PATCH that carries `scoring` is counted. Rewriting `content`, `provenance` + or `access_policy` is unaffected. `{"scoring": null}` changes nothing, so it + costs nothing. +- The budget is per cell, so one busy cell cannot use up another's. +- Refused edits are `429 RATE_LIMITED`, with `Retry-After` in seconds and the same + number in `error.details.retry_after_seconds`. The number is how long until the + oldest allowed edit leaves the window - a refused attempt is not recorded, so + retrying early does not push your own deadline back. +- Access is checked first: a caller who may not write the cell gets `403` and + learns nothing about the remaining budget. +- Default: 5 edits per cell per hour. `AMP_SCORING_PATCH_LIMIT` and + `AMP_SCORING_PATCH_WINDOW_SECONDS` change it; `AMP_SCORING_PATCH_LIMIT=0` + disables it, and [`GET /spec`](#get-spec) then reports `null`. + +The counters live in the server process, so two processes over one storage backend +keep two budgets. + +--- + ## DELETE /memories/{memory_id} Soft-deletes a memory cell by setting `lifecycle.status` to `"deleted"`. The cell is retained in storage and will not appear in search results, but can still be retrieved directly by ID. diff --git a/server/amp_server/errors.py b/server/amp_server/errors.py index 8ccbee5..2915ee3 100644 --- a/server/amp_server/errors.py +++ b/server/amp_server/errors.py @@ -26,11 +26,15 @@ def __init__( code: str, message: str, details: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, ) -> None: self.status_code = status_code self.code = code self.message = message self.details = details or {} + # Part of the error for a caller who can act on it: a rate limit that + # does not say how long to wait leaves the client to guess. + self.headers = headers or {} super().__init__(message) def to_response(self) -> dict[str, Any]: @@ -71,6 +75,21 @@ def unauthenticated() -> AMPError: ) +def rate_limited(retry_after_seconds: int) -> AMPError: + """Too many scoring edits on one cell (RFC §5: decay-score manipulation). + + `429` with `Retry-After`, because the caller can act on the number: it is + exactly how long until the oldest allowed edit ages out of the window. + """ + return AMPError( + 429, + "RATE_LIMITED", + f"too many scoring updates for this cell; retry in {retry_after_seconds}s", + details={"retry_after_seconds": retry_after_seconds}, + headers={"Retry-After": str(retry_after_seconds)}, + ) + + def admin_disabled() -> AMPError: """Manual lifecycle runs are opt-in: no token configured means no access.""" return AMPError( diff --git a/server/amp_server/main.py b/server/amp_server/main.py index d706247..7200dfd 100644 --- a/server/amp_server/main.py +++ b/server/amp_server/main.py @@ -21,6 +21,7 @@ admin_disabled, invalid_transition, missing_agent_id, + rate_limited, unauthenticated, ) from amp_server.lifecycle import LifecycleEngine @@ -41,6 +42,11 @@ SearchResponse, ) from amp_server.paging import DEFAULT_PAGE_SIZE, MAX_PAGE_SIZE, readable_page +from amp_server.ratelimit import ( + ScoringPatchLimiter, + limit_from_env, + patch_touches_scoring, +) from amp_server.scheduler import lifecycle_loop, run_lifecycle, settings_from_env from amp_server.storage.base import ( InvalidTransitionError, @@ -75,6 +81,10 @@ def configure_logging() -> None: # trusted as the spec's binding describes. Set from the lifespan, so a key store # that cannot be read stops the server rather than turning into a 401 later. _api_key_store: ApiKeyStore | None = None +# Resolved here, like the lifecycle settings, so a test can swap either without +# reaching into the environment. The limiter is the stateful half. +_scoring_limit = limit_from_env() +_scoring_limiter = ScoringPatchLimiter(_scoring_limit) def _build_storage() -> StorageAdapter: @@ -162,6 +172,7 @@ def get_lifecycle() -> LifecycleEngine: @asynccontextmanager async def lifespan(app: FastAPI): global _storage, _lifecycle, _lifecycle_settings, _api_key_store + global _scoring_limit, _scoring_limiter configure_logging() # Built here, not at import: a misconfigured backend must stop the server # from starting rather than surface as poor search results later. @@ -171,6 +182,17 @@ async def lifespan(app: FastAPI): # fall back to trusting an unverified header. _api_key_store = store_from_env() + # Read again here for the same reason the lifecycle settings are: a test (or + # an embedder) can set AMP_SCORING_PATCH_LIMIT before starting the app. + _scoring_limit = limit_from_env() + _scoring_limiter = ScoringPatchLimiter(_scoring_limit) + if _scoring_limit.enabled: + logger.info( + "Scoring PATCH limit: %d per %ds per cell", + _scoring_limit.max_patches, + _scoring_limit.window_seconds, + ) + # Read settings here, not only at import: a test (or an embedder) can set # AMP_LIFECYCLE_* before starting the app and expect it to take effect. _lifecycle_settings = settings_from_env() @@ -227,7 +249,11 @@ async def _amp_error_handler(request: Request, exc: AMPError) -> JSONResponse: `{"error": {"code", "message", "details"}}` body instead of each route building its own, and so route handlers can stay annotated `-> dict`. """ - return JSONResponse(status_code=exc.status_code, content=exc.to_response()) + return JSONResponse( + status_code=exc.status_code, + content=exc.to_response(), + headers=exc.headers or None, + ) # The error envelope is part of the protocol, so it belongs in the contract @@ -247,6 +273,13 @@ async def _amp_error_handler(request: Request, exc: AMPError) -> JSONResponse: "(only when AMP_API_KEYS_FILE is configured)", } } +_RATE_LIMITED: dict[int | str, dict[str, Any]] = { + 429: { + "model": ErrorResponse, + "description": "Too many scoring updates for this cell; see Retry-After " + "and GET /spec's scoring_patch_limit", + } +} _ACCESS_DENIED: dict[int | str, dict[str, Any]] = { 403: { "model": ErrorResponse, @@ -288,6 +321,7 @@ async def spec() -> dict[str, Any]: # before it makes a call that would fail with 401. "api_keys_required": _api_key_store is not None, "max_page_size": MAX_PAGE_SIZE, + "scoring_patch_limit": _scoring_limit.to_json(), "embedding": get_storage().embedding, "max_cell_size_bytes": MAX_CELL_SIZE_BYTES, "retention_days": get_storage().retention_days, @@ -452,6 +486,7 @@ async def get_memory( | _ACCESS_DENIED | _INVALID_TRANSITION | _CELL_TOO_LARGE + | _RATE_LIMITED ), ) async def update_memory( @@ -475,6 +510,17 @@ async def update_memory( if not check_write_access(cell, x_amp_agent_id): raise access_denied() + if patch_touches_scoring(body): + # RFC §5: a caller looping on `scoring` can hold a cell active past its + # window or force a competitor into archive. Checked after the access + # rules, so a caller who may not write this cell is told that, and told + # nothing about how much budget is left. An allowed edit is recorded + # here, before the write, so a write that then fails does not buy the + # caller a free retry. + wait = _scoring_limiter.check_and_record(memory_id) + if wait: + raise rate_limited(wait) + try: updated = await storage.update(memory_id, body) except InvalidTransitionError as exc: @@ -516,6 +562,10 @@ async def delete_memory( except InvalidTransitionError as exc: raise invalid_transition(exc.message) from exc + # The cell is gone from the API's point of view; its scoring budget has + # nothing left to protect, and holding the counters would keep an id alive + # in memory for a cell no caller can reach. + _scoring_limiter.forget(memory_id) return Response(status_code=204) diff --git a/server/amp_server/ratelimit.py b/server/amp_server/ratelimit.py new file mode 100644 index 0000000..44d5088 --- /dev/null +++ b/server/amp_server/ratelimit.py @@ -0,0 +1,178 @@ +"""A per-cell budget on rewriting `scoring` (RFC-AMP-001 §5). + +The RFC names decay-score manipulation as a threat: a caller PATCHing `scoring` +in a loop can hold a cell `active` past its relevance window, or drive a competing +memory into archive. Its first mitigation is structural and already in place - a +scoring PATCH only takes effect on the next `LifecycleEngine.process_all`, so the +engine's cadence bounds how fast manipulation can land. Its second is this, asking +implementations to rate-limit scoring PATCH frequency per cell. Nothing did. + +Two choices worth stating: + +- **Only a PATCH carrying `scoring` is counted.** Rewriting text, metadata or the + access policy is not what the threat describes, and counting those would refuse + ordinary edits to cover an attack they have nothing to do with. +- **A refused attempt is not recorded.** Counting refusals would let a caller + extend its own lockout by hammering, and would make `Retry-After` grow while the + caller waits. Not recording means the wait shrinks as the oldest allowed edit + ages out, which is the behaviour the number promises. + +The counters live in this process. Two server processes over one Postgres keep two +budgets - worth knowing, and better than a shared counter that would need its own +storage and its own consistency story. The threat is one caller in a loop, and one +process sees that caller's whole stream. +""" + +from __future__ import annotations + +import logging +import math +import os +import time +from collections import OrderedDict, deque +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from amp_server.models import MemoryCellUpdate + +logger = logging.getLogger(__name__) + +#: Scoring edits allowed per cell per window, when nothing else is configured. +#: Five is generous for a person re-scoring a memory and useless for a loop; the +#: lifecycle engine runs hourly by default, so this is five edits per cycle. +DEFAULT_MAX_PATCHES = 5 +DEFAULT_WINDOW_SECONDS = 3600 + +#: How many cells are tracked at once. Counters for cells nobody edits are +#: dropped as their window empties; this caps the rest. +DEFAULT_MAX_TRACKED_CELLS = 10_000 + + +@dataclass(frozen=True) +class ScoringPatchLimit: + """How often one cell's `scoring` may be rewritten. + + `max_patches` of 0 disables the limit, which is the honest way to express + "this deployment does not want it" rather than a very large number that + reads like a policy. + """ + + max_patches: int = DEFAULT_MAX_PATCHES + window_seconds: int = DEFAULT_WINDOW_SECONDS + + @property + def enabled(self) -> bool: + return self.max_patches > 0 and self.window_seconds > 0 + + def to_json(self) -> dict[str, Any] | None: + """What `GET /spec` reports, or None when the limit is off.""" + if not self.enabled: + return None + return { + "max_patches": self.max_patches, + "window_seconds": self.window_seconds, + } + + +def limit_from_env() -> ScoringPatchLimit: + """Read the limit from the environment, falling back rather than failing. + + A misread number here degrades a mitigation, it does not corrupt data, so + this warns and uses the default instead of stopping the server - unlike the + embedding provider and the key store, where a wrong value is worse than no + server at all. + """ + return ScoringPatchLimit( + max_patches=_int_env("AMP_SCORING_PATCH_LIMIT", DEFAULT_MAX_PATCHES), + window_seconds=_int_env( + "AMP_SCORING_PATCH_WINDOW_SECONDS", DEFAULT_WINDOW_SECONDS + ), + ) + + +def _int_env(name: str, default: int) -> int: + raw = os.environ.get(name) + if raw is None or raw.strip() == "": + return default + try: + value = int(raw) + except ValueError: + logger.warning("%s=%r is not an integer; using %d", name, raw, default) + return default + if value < 0: + logger.warning("%s must not be negative; using %d", name, default) + return default + return value + + +class ScoringPatchLimiter: + """Counts scoring edits per cell, in memory, over a sliding window.""" + + def __init__( + self, + limit: ScoringPatchLimit, + *, + clock: Callable[[], float] = time.monotonic, + max_tracked_cells: int = DEFAULT_MAX_TRACKED_CELLS, + ) -> None: + self.limit = limit + self._clock = clock + self._max_tracked = max_tracked_cells + self._seen: OrderedDict[str, deque[float]] = OrderedDict() + + def check_and_record(self, cell_id: str) -> int: + """Count one attempt and return the seconds to wait; 0 when allowed. + + One call rather than `check` plus `record`, so a caller cannot pass the + check and forget to spend the budget. + """ + if not self.limit.enabled: + return 0 + + now = self._clock() + stamps = self._seen.get(cell_id) + if stamps is None: + stamps = deque() + self._seen[cell_id] = stamps + else: + self._seen.move_to_end(cell_id) + self._prune(stamps, now) + + if len(stamps) >= self.limit.max_patches: + return max(1, math.ceil(self.limit.window_seconds - (now - stamps[0]))) + + stamps.append(now) + self._evict_if_crowded() + return 0 + + def forget(self, cell_id: str) -> None: + """Drop a cell's counters, e.g. once it has been purged.""" + self._seen.pop(cell_id, None) + + def tracked_cells(self) -> int: + """How many cells are being counted, for tests and diagnostics.""" + return len(self._seen) + + def _prune(self, stamps: deque[float], now: float) -> None: + while stamps and now - stamps[0] >= self.limit.window_seconds: + stamps.popleft() + + def _evict_if_crowded(self) -> None: + """Forget the least recently touched cells once the cap is reached. + + Bounded memory matters more than perfect accounting for cells that were + edited once, long ago: evicting the oldest entry can only ever make the + limit more permissive, never refuse an edit that should be allowed. + """ + while len(self._seen) > self._max_tracked: + self._seen.popitem(last=False) + + +def patch_touches_scoring(body: MemoryCellUpdate) -> bool: + """Whether this PATCH actually rewrites `scoring`. + + `{"scoring": null}` sends nothing to change, so it costs nothing: the budget + is spent by writes, not by requests. + """ + return body.scoring is not None diff --git a/server/tests/test_ratelimit.py b/server/tests/test_ratelimit.py new file mode 100644 index 0000000..dee0e9d --- /dev/null +++ b/server/tests/test_ratelimit.py @@ -0,0 +1,348 @@ +"""Rate-limiting scoring edits, per cell (RFC-AMP-001 §5). + +The RFC lists decay-score manipulation as a threat and asks implementations to +rate-limit how often one cell's `scoring` may be rewritten: a caller looping on it +can hold a cell `active` past its relevance window, or drive a competing memory +into archive. Two halves matter here - that a loop is refused, and that ordinary +editing is not caught by the refusal. A limit that refuses both is not a +mitigation, it is an outage. +""" + +from __future__ import annotations + +import uuid + +import pytest +from conftest import make_cell +from httpx import ASGITransport, AsyncClient, Response + +from amp_server.models import LifecycleStatus, MemoryCellUpdate, MemoryScoring +from amp_server.ratelimit import ( + DEFAULT_MAX_PATCHES, + DEFAULT_WINDOW_SECONDS, + ScoringPatchLimit, + ScoringPatchLimiter, + limit_from_env, + patch_touches_scoring, +) + +_AGENT = "agent-ratelimit" + + +class FakeClock: + """A clock the test moves, so a window costs no wall time.""" + + def __init__(self, start: float = 1_000.0) -> None: + self.now = start + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +def _install_app( + limit: ScoringPatchLimit | None = None, clock: FakeClock | None = None +) -> ScoringPatchLimiter: + """Fresh storage plus a limiter the test controls.""" + import amp_server.main as main_mod + from amp_server.lifecycle import LifecycleEngine + from amp_server.storage.chroma import ChromaAdapter + + main_mod._storage = ChromaAdapter(collection_name=f"test_{uuid.uuid4().hex[:12]}") + main_mod._lifecycle = LifecycleEngine(main_mod._storage) + main_mod._api_key_store = None + + resolved = limit or ScoringPatchLimit() + limiter = ScoringPatchLimiter(resolved, clock=clock or FakeClock()) + main_mod._scoring_limit = resolved + main_mod._scoring_limiter = limiter + return limiter + + +async def _client() -> AsyncClient: + import amp_server.main as main_mod + + return AsyncClient( + transport=ASGITransport(app=main_mod.app), base_url="http://test" + ) + + +async def _store_cell(text: str = "scored cell") -> str: + import amp_server.main as main_mod + + cell = make_cell(owner_id=_AGENT, created_by=_AGENT, text=text) + await main_mod.get_storage().save(cell) + return cell.id + + +def _scoring(importance: float = 0.9) -> dict: + return {"scoring": {"importance": importance, "confidence": 1.0, "decay_rate": 0.0}} + + +def _headers() -> dict[str, str]: + return {"X-AMP-Agent-ID": _AGENT} + + +async def _patch(memory_id: str, body: dict) -> Response: + async with await _client() as client: + return await client.patch( + f"/amp/v1/memories/{memory_id}", json=body, headers=_headers() + ) + + +# --------------------------------------------------------------------------- +# The rule itself +# --------------------------------------------------------------------------- + + +def test_a_null_scoring_field_costs_nothing(): + """The budget is spent by writes, and `{"scoring": null}` writes nothing.""" + assert patch_touches_scoring( + MemoryCellUpdate.model_validate({"scoring": None}) + ) is (False) + assert ( + patch_touches_scoring( + MemoryCellUpdate.model_validate({"content": {"text": "x"}}) + ) + is False + ) + assert not patch_touches_scoring(MemoryCellUpdate.model_validate({"scoring": None})) + assert not patch_touches_scoring( + MemoryCellUpdate.model_validate({"content": {"text": "x"}}) + ) + assert patch_touches_scoring(MemoryCellUpdate(scoring=MemoryScoring())) + + +def test_the_limit_refuses_past_its_budget_and_says_how_long(): + clock = FakeClock() + limiter = ScoringPatchLimiter( + ScoringPatchLimit(max_patches=2, window_seconds=60), clock=clock + ) + + assert limiter.check_and_record("mem_a") == 0 + assert limiter.check_and_record("mem_a") == 0 + assert limiter.check_and_record("mem_a") == 60 # the first edit ages out then + + +def test_a_refusal_does_not_extend_the_lockout(): + """Otherwise hammering would push the caller's own deadline back.""" + clock = FakeClock() + limiter = ScoringPatchLimiter( + ScoringPatchLimit(max_patches=1, window_seconds=60), clock=clock + ) + limiter.check_and_record("mem_a") + + clock.advance(10) + first_refusal = limiter.check_and_record("mem_a") + clock.advance(10) + second_refusal = limiter.check_and_record("mem_a") + + assert first_refusal == 50 + assert second_refusal == 40, "the wait must shrink, not grow" + + +def test_the_window_is_per_cell(): + clock = FakeClock() + limiter = ScoringPatchLimiter( + ScoringPatchLimit(max_patches=1, window_seconds=60), clock=clock + ) + limiter.check_and_record("mem_a") + + assert limiter.check_and_record("mem_a") == 60 + assert limiter.check_and_record("mem_b") == 0 + + +def test_the_budget_returns_once_the_window_passes(): + clock = FakeClock() + limiter = ScoringPatchLimiter( + ScoringPatchLimit(max_patches=1, window_seconds=60), clock=clock + ) + limiter.check_and_record("mem_a") + + clock.advance(59) + assert limiter.check_and_record("mem_a") == 1 + clock.advance(1) + assert limiter.check_and_record("mem_a") == 0 + + +def test_cells_are_forgotten_after_a_purge(): + limiter = ScoringPatchLimiter(ScoringPatchLimit(max_patches=1, window_seconds=60)) + + limiter.check_and_record("mem_a") + assert limiter.tracked_cells() == 1 + + limiter.forget("mem_a") + assert limiter.tracked_cells() == 0 + + +def test_tracking_is_bounded(): + """An unbounded counter map is a slow leak in a long-running server.""" + limiter = ScoringPatchLimiter( + ScoringPatchLimit(max_patches=1, window_seconds=60), max_tracked_cells=3 + ) + + for index in range(10): + limiter.check_and_record(f"mem_{index}") + + assert limiter.tracked_cells() == 3 + + +def test_a_disabled_limit_never_refuses(): + clock = FakeClock() + limiter = ScoringPatchLimiter( + ScoringPatchLimit(max_patches=0, window_seconds=60), clock=clock + ) + + for _ in range(50): + assert limiter.check_and_record("mem_a") == 0 + assert ScoringPatchLimit(max_patches=0).to_json() is None + + +def test_the_settings_default_to_the_documented_numbers(monkeypatch): + for var in ("AMP_SCORING_PATCH_LIMIT", "AMP_SCORING_PATCH_WINDOW_SECONDS"): + monkeypatch.delenv(var, raising=False) + + limit = limit_from_env() + assert limit.max_patches == DEFAULT_MAX_PATCHES + assert limit.window_seconds == DEFAULT_WINDOW_SECONDS + assert limit.enabled is True + + +def test_a_bad_setting_falls_back_rather_than_stopping_the_server(monkeypatch): + """A wrong number here degrades a mitigation; it does not corrupt data.""" + monkeypatch.setenv("AMP_SCORING_PATCH_LIMIT", "not-a-number") + assert limit_from_env().max_patches == DEFAULT_MAX_PATCHES + + monkeypatch.setenv("AMP_SCORING_PATCH_LIMIT", "-1") + assert limit_from_env().max_patches == DEFAULT_MAX_PATCHES + + monkeypatch.setenv("AMP_SCORING_PATCH_LIMIT", "0") + assert limit_from_env().enabled is False + + +# --------------------------------------------------------------------------- +# Over HTTP +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_loop_on_scoring_is_refused_with_a_retry_after(): + _install_app(ScoringPatchLimit(max_patches=2, window_seconds=3600)) + cell_id = await _store_cell() + + first = await _patch(cell_id, _scoring(0.9)) + second = await _patch(cell_id, _scoring(0.8)) + third = await _patch(cell_id, _scoring(0.7)) + + assert [first.status_code, second.status_code] == [200, 200] + assert third.status_code == 429 + assert third.json()["error"]["code"] == "RATE_LIMITED" + assert third.headers["Retry-After"] == "3600" + # The refused write did not land. + assert third.json()["error"]["details"]["retry_after_seconds"] == 3600 + + +@pytest.mark.asyncio +async def test_ordinary_edits_are_not_counted(): + """The limit is about `scoring`; refusing text edits would be an outage.""" + _install_app(ScoringPatchLimit(max_patches=1, window_seconds=3600)) + cell_id = await _store_cell() + + assert (await _patch(cell_id, _scoring(0.9))).status_code == 200 + for index in range(5): + response = await _patch(cell_id, {"content": {"text": f"rewritten {index}"}}) + assert response.status_code == 200, f"edit {index} was counted as scoring" + + assert (await _patch(cell_id, _scoring(0.5))).status_code == 429 + + +@pytest.mark.asyncio +async def test_one_cell_exhausting_its_budget_does_not_block_another(): + _install_app(ScoringPatchLimit(max_patches=1, window_seconds=3600)) + busy = await _store_cell("busy") + other = await _store_cell("other") + + assert (await _patch(busy, _scoring(0.9))).status_code == 200 + assert (await _patch(busy, _scoring(0.8))).status_code == 429 + assert (await _patch(other, _scoring(0.8))).status_code == 200 + + +@pytest.mark.asyncio +async def test_the_budget_returns_when_the_window_passes_over_http(): + clock = FakeClock() + _install_app(ScoringPatchLimit(max_patches=1, window_seconds=60), clock=clock) + cell_id = await _store_cell() + + assert (await _patch(cell_id, _scoring(0.9))).status_code == 200 + assert (await _patch(cell_id, _scoring(0.8))).status_code == 429 + + clock.advance(60) + assert (await _patch(cell_id, _scoring(0.7))).status_code == 200 + + +@pytest.mark.asyncio +async def test_a_caller_without_write_access_learns_nothing_about_the_budget(): + """Access is decided first, so a stranger cannot probe the limit.""" + _install_app(ScoringPatchLimit(max_patches=1, window_seconds=3600)) + import amp_server.main as main_mod + + cell = make_cell(owner_id="user-someone-else", created_by="agent-someone-else") + await main_mod.get_storage().save(cell) + + response = await _patch(cell.id, _scoring(0.9)) + + assert response.status_code == 403 + assert response.json()["error"]["code"] == "ACCESS_DENIED" + + +@pytest.mark.asyncio +async def test_deleting_a_cell_forgets_its_budget(): + _install_app(ScoringPatchLimit(max_patches=1, window_seconds=3600)) + import amp_server.main as main_mod + + cell = make_cell( + owner_id=_AGENT, created_by=_AGENT, status=LifecycleStatus.ARCHIVED + ) + await main_mod.get_storage().save(cell) + assert (await _patch(cell.id, _scoring(0.9))).status_code == 200 + + async with await _client() as client: + deleted = await client.delete(f"/amp/v1/memories/{cell.id}", headers=_headers()) + + assert deleted.status_code == 204 + assert main_mod._scoring_limiter.tracked_cells() == 0 + + +@pytest.mark.asyncio +async def test_spec_advertises_the_budget_it_enforces(): + _install_app(ScoringPatchLimit(max_patches=2, window_seconds=900)) + cell_id = await _store_cell() + + async with await _client() as client: + advertised = (await client.get("/amp/v1/spec")).json()["capabilities"][ + "scoring_patch_limit" + ] + + assert advertised == {"max_patches": 2, "window_seconds": 900} + + # And the advertised number is the one the route uses. + assert (await _patch(cell_id, _scoring(0.9))).status_code == 200 + assert (await _patch(cell_id, _scoring(0.8))).status_code == 200 + assert (await _patch(cell_id, _scoring(0.7))).status_code == 429 + + +@pytest.mark.asyncio +async def test_spec_says_null_when_the_limit_is_off(): + _install_app(ScoringPatchLimit(max_patches=0)) + cell_id = await _store_cell() + + async with await _client() as client: + advertised = (await client.get("/amp/v1/spec")).json()["capabilities"][ + "scoring_patch_limit" + ] + + assert advertised is None + for _ in range(10): + assert (await _patch(cell_id, _scoring(0.9))).status_code == 200 diff --git a/spec/v0.1.0/openapi.json b/spec/v0.1.0/openapi.json index 953c15e..a73b0b0 100644 --- a/spec/v0.1.0/openapi.json +++ b/spec/v0.1.0/openapi.json @@ -1364,6 +1364,16 @@ } }, "description": "Validation Error" + }, + "429": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "Too many scoring updates for this cell; see Retry-After and GET /spec's scoring_patch_limit" } }, "summary": "Update Memory"