From 3e0b6584317445d32f96470ef712d1e61d5d62bf Mon Sep 17 00:00:00 2001 From: lr90 Date: Tue, 20 Oct 2026 00:44:07 +0800 Subject: [PATCH] fix(adapter): separate layered timeout and request budgets --- agent_memory_leaderboard/README.md | 36 +++- agent_memory_leaderboard/src/app.py | 17 +- agent_memory_leaderboard/src/benchmark.py | 4 +- agent_memory_leaderboard/src/budget.py | 41 +++++ agent_memory_leaderboard/src/config.py | 18 +- agent_memory_leaderboard/src/diagnostics.py | 6 + agent_memory_leaderboard/src/journal.py | 22 ++- agent_memory_leaderboard/src/memoria.py | 12 +- agent_memory_leaderboard/src/probes.py | 14 +- agent_memory_leaderboard/src/scheduling.py | 8 +- agent_memory_leaderboard/src/service.py | 11 +- agent_memory_leaderboard/src/smoke.py | 2 +- .../tests/test_request_budget.py | 174 ++++++++++++++++++ 13 files changed, 341 insertions(+), 24 deletions(-) create mode 100644 agent_memory_leaderboard/src/budget.py create mode 100644 agent_memory_leaderboard/tests/test_request_budget.py diff --git a/agent_memory_leaderboard/README.md b/agent_memory_leaderboard/README.md index 03f06ec..34302e7 100644 --- a/agent_memory_leaderboard/README.md +++ b/agent_memory_leaderboard/README.md @@ -195,7 +195,7 @@ never held across a Memoria HTTP request. The existing nonblocking per-request lock and its HTTP 409 behavior are unchanged. `ADAPTER_JOURNAL_QUEUE_TIMEOUT_SECONDS` sets the slot acquisition budget -(including both waits), default `5`, finite and in `(0, 120]`. It is independent +(including both waits), default `30`, finite and in `(0, 120]`. It is independent of the application request limit: raising capacity does not raise durable-write throughput, so choose the waiting budget explicitly. An exhausted wait returns a safe HTTP 503 and records `journal_queue_timeout`; lock-file I/O @@ -299,7 +299,7 @@ payload. It does not retry Search 409/524 or malformed successful responses. ## Write visibility budget -`MEMORIA_VISIBILITY_TIMEOUT_SECONDS` defaults to `10`; it must be finite and in +`MEMORIA_VISIBILITY_TIMEOUT_SECONDS` defaults to `60`; it must be finite and in `(0, 120]`. After a write has a known memory ID, Add waits for a valid point read and retrieval of that exact write. Existing known-ID inflight rows (including legacy rows) use the same path. A valid null read or valid retrieval result that @@ -570,3 +570,35 @@ See [evaluation and diagnostics](benchmarks/EVALUATION.md) for label-free source attribution, optional Answer/Eval, model settings, private caches, checkpoints, source-time regression cases and bounded capacity probes. Retrieval remains the default and does not require answer/judge model credentials. + + +### Layered timeout budgets + +Candidate defaults are journal-slot queue 30s, backend connection/pool 30s, +backend read/write I/O 120s, visibility polling 60s, and admission/backend queue +120s. `ADAPTER_JOURNAL_QUEUE_TIMEOUT_SECONDS` remains bounded to 120s. Raising +its default from 5s reduces premature queue failures without increasing writer +throughput. SQLite's own busy timeout remains 5s; probes retain their separate +short read/health budgets. + +`MEMORIA_CONNECT_TIMEOUT_SECONDS` configures connection and HTTP pool waiting; +`MEMORIA_TIMEOUT_SECONDS` configures individual read/write I/O waits. +`MEMORIA_VISIBILITY_TIMEOUT_SECONDS` covers point-read plus search verification, +and pending visibility never repeats a backend write. + +`ADAPTER_REQUEST_TIMEOUT_SECONDS` defaults to 1700s and accepts at most 1740s. +It establishes one cooperative budget across admission, journal waits, all Add +chunks and backend calls. Later phases use the remaining budget, and expiry +returns retryable 504. It is not a hard HTTP stopwatch: an in-flight synchronous +I/O or SQLite commit unwinds safely before its thread and admission permit are +released. A late committed Add remains recoverable using the original request +ID, even if the response is a timeout. The external API I/O window, probe and +shutdown settings are independent. The local smoke, formal benchmark and probe +clients use a 1800s I/O window (or the selected probe window) with a 30s maximum +connection/pool wait; answer/judge model-call budgets remain unchanged. + +QA and AWS charts must pass the same explicit timeout values and deploy an +image containing these settings. Unknown environment variables are silently +ignored by older adapter images. These budgets are candidates to validate with +64-client mixed and multi-message loads; no deployment or throughput guarantee +is implied by changing them. diff --git a/agent_memory_leaderboard/src/app.py b/agent_memory_leaderboard/src/app.py index ec79349..425bee5 100644 --- a/agent_memory_leaderboard/src/app.py +++ b/agent_memory_leaderboard/src/app.py @@ -8,6 +8,7 @@ from fastapi.responses import JSONResponse from starlette.exceptions import HTTPException +from .budget import request_budget from .config import Settings from .errors import AdapterError from .journal import Journal @@ -127,9 +128,12 @@ async def validation_error(request: Request, error: RequestValidationError): @app.post("/add", response_model=AddResponse, dependencies=[Depends(authenticate)]) async def add(request: AddRequest): - async with admission.slot(): - # Keep admission until the write thread finishes and records its outcome. - return await anyio.to_thread.run_sync(adapter.add, request, limiter=business_threads) + with request_budget(settings.request_timeout_seconds): + async with admission.slot(): + # Keep admission until the write thread finishes and records its outcome. + return await anyio.to_thread.run_sync( + adapter.add, request, limiter=business_threads + ) @app.post( "/search", @@ -138,7 +142,10 @@ async def add(request: AddRequest): dependencies=[Depends(authenticate)], ) async def search(request: SearchRequest): - async with admission.slot(): - return await anyio.to_thread.run_sync(adapter.search, request, limiter=business_threads) + with request_budget(settings.request_timeout_seconds): + async with admission.slot(): + return await anyio.to_thread.run_sync( + adapter.search, request, limiter=business_threads + ) return app diff --git a/agent_memory_leaderboard/src/benchmark.py b/agent_memory_leaderboard/src/benchmark.py index 3dd4925..09857a7 100644 --- a/agent_memory_leaderboard/src/benchmark.py +++ b/agent_memory_leaderboard/src/benchmark.py @@ -200,7 +200,9 @@ def call(http, path, body, identity): persist() try: - with httpx.Client(timeout=1800, trust_env=False) as http: + with httpx.Client( + timeout=httpx.Timeout(1800, connect=30, pool=30), trust_env=False + ) as http: http.get(args.adapter_url.rstrip("/") + "/health").raise_for_status() for conversation in selected: sample = conversation["sample_id"] diff --git a/agent_memory_leaderboard/src/budget.py b/agent_memory_leaderboard/src/budget.py new file mode 100644 index 0000000..38d125d --- /dev/null +++ b/agent_memory_leaderboard/src/budget.py @@ -0,0 +1,41 @@ +"""Cooperative request budget shared by admission, journal and backend work. + +An in-flight synchronous I/O or commit is allowed to unwind safely. This is not +an asynchronous hard kill and does not claim a strict HTTP wall-clock deadline. +""" + +import time +from contextlib import contextmanager +from contextvars import ContextVar + +from .errors import AdapterError + +DEADLINE: ContextVar[float | None] = ContextVar("adapter_request_deadline", default=None) + + +class RequestBudgetExceeded(AdapterError): + def __init__(self): + super().__init__( + 504, "Adapter request budget exhausted; retry the same request", retry_after="2" + ) + + +def remaining(limit: float) -> float: + deadline = DEADLINE.get() + if deadline is None: + return limit + available = deadline - time.monotonic() + if available <= 0: + raise RequestBudgetExceeded() + return min(limit, available) + + +@contextmanager +def request_budget(seconds: float): + parent = DEADLINE.get() + deadline = time.monotonic() + seconds + token = DEADLINE.set(min(parent, deadline) if parent is not None else deadline) + try: + yield + finally: + DEADLINE.reset(token) diff --git a/agent_memory_leaderboard/src/config.py b/agent_memory_leaderboard/src/config.py index de7622a..b4b5b6b 100644 --- a/agent_memory_leaderboard/src/config.py +++ b/agent_memory_leaderboard/src/config.py @@ -13,14 +13,16 @@ class Settings: state_dir: Path = Path("var") namespace: str = "agent-memory-leaderboard" timeout_seconds: float = 120 - visibility_timeout_seconds: float = field(default=10, kw_only=True) + connect_timeout_seconds: float = field(default=30, kw_only=True) + request_timeout_seconds: float = field(default=1700, kw_only=True) + visibility_timeout_seconds: float = field(default=60, kw_only=True) backend_max_concurrency: int = 0 backend_requests_per_second: float = 0 backend_queue_timeout_seconds: float = 120 max_active_requests: int = 40 journal_log_all_operations: bool = False journal_slow_ms: float = 100 - journal_queue_timeout_seconds: float = 5 + journal_queue_timeout_seconds: float = 30 def __post_init__(self) -> None: url = urlsplit(self.memoria_url) @@ -35,6 +37,12 @@ def __post_init__(self) -> None: raise ValueError("MEMORIA_URL must be an HTTP(S) base URL without credentials") if not self.memoria_master_key or not self.namespace.strip(): raise ValueError("MEMORIA_MASTER_KEY and MEMORIA_NAMESPACE must be nonempty") + for name, value, maximum in [ + ("MEMORIA_CONNECT_TIMEOUT_SECONDS", self.connect_timeout_seconds, 120), + ("ADAPTER_REQUEST_TIMEOUT_SECONDS", self.request_timeout_seconds, 1740), + ]: + if not math.isfinite(value) or not 0 < value <= maximum: + raise ValueError(f"{name} must be in (0, {maximum}]") if not math.isfinite(self.timeout_seconds) or self.timeout_seconds <= 0: raise ValueError("MEMORIA_TIMEOUT_SECONDS must be positive and finite") if ( @@ -78,7 +86,9 @@ def from_env(cls) -> "Settings": state_dir=Path(os.getenv("ADAPTER_STATE_DIR", "var")), namespace=os.getenv("MEMORIA_NAMESPACE", "agent-memory-leaderboard"), timeout_seconds=float(os.getenv("MEMORIA_TIMEOUT_SECONDS", "120")), - visibility_timeout_seconds=float(os.getenv("MEMORIA_VISIBILITY_TIMEOUT_SECONDS", "10")), + connect_timeout_seconds=float(os.getenv("MEMORIA_CONNECT_TIMEOUT_SECONDS", "30")), + request_timeout_seconds=float(os.getenv("ADAPTER_REQUEST_TIMEOUT_SECONDS", "1700")), + visibility_timeout_seconds=float(os.getenv("MEMORIA_VISIBILITY_TIMEOUT_SECONDS", "60")), backend_max_concurrency=int(os.getenv("MEMORIA_MAX_CONCURRENCY", "0")), backend_requests_per_second=float(os.getenv("MEMORIA_REQUESTS_PER_SECOND", "0")), backend_queue_timeout_seconds=float(os.getenv("MEMORIA_QUEUE_TIMEOUT_SECONDS", "120")), @@ -86,6 +96,6 @@ def from_env(cls) -> "Settings": journal_log_all_operations=journal_logging in {"true", "1"}, journal_slow_ms=float(os.getenv("ADAPTER_JOURNAL_SLOW_MS", "100")), journal_queue_timeout_seconds=float( - os.getenv("ADAPTER_JOURNAL_QUEUE_TIMEOUT_SECONDS", "5") + os.getenv("ADAPTER_JOURNAL_QUEUE_TIMEOUT_SECONDS", "30") ), ) diff --git a/agent_memory_leaderboard/src/diagnostics.py b/agent_memory_leaderboard/src/diagnostics.py index d49fd95..5e4045d 100644 --- a/agent_memory_leaderboard/src/diagnostics.py +++ b/agent_memory_leaderboard/src/diagnostics.py @@ -8,6 +8,12 @@ def adapter_reason_code(response): """Recognize fixed adapter reasons without persisting response text or values.""" reasons = { + ( + 504, + "Adapter request budget exhausted; retry the same request", + ): "adapter_request_timeout", + (503, "Idempotency journal is busy; retry Add"): "journal_queue_timeout", + (503, "Memoria write visibility timed out; retry Add"): "write_visibility_timeout", (503, "Prior write outcome is unknown; retry the same Add later"): "write_outcome_unknown", ( 503, diff --git a/agent_memory_leaderboard/src/journal.py b/agent_memory_leaderboard/src/journal.py index a667666..4a44096 100644 --- a/agent_memory_leaderboard/src/journal.py +++ b/agent_memory_leaderboard/src/journal.py @@ -8,6 +8,8 @@ from contextlib import contextmanager from pathlib import Path +from .budget import RequestBudgetExceeded +from .budget import remaining as budget_remaining from .errors import AdapterError from .journal_metrics import JournalMetrics @@ -21,6 +23,7 @@ def __init__(self, db, timing): self.timing = timing def execute(self, statement, parameters=()): + budget_remaining(float("inf")) return self.timing.run("sql", self.db.execute, statement, parameters) @@ -31,7 +34,7 @@ def __init__( *, log_all_operations=False, slow_ms=100, - queue_timeout_seconds=5.0, + queue_timeout_seconds=30.0, ): if not math.isfinite(queue_timeout_seconds) or not 0 < queue_timeout_seconds <= 120: raise ValueError("Journal queue timeout must be in (0, 120]") @@ -71,8 +74,9 @@ def __init__( db.execute("ALTER TABLE writes ADD COLUMN protocol TEXT NOT NULL DEFAULT 'legacy'") def _acquire_connection_slot(self): - deadline = time.monotonic() + self.queue_timeout_seconds - if not self._connection_mutex.acquire(timeout=self.queue_timeout_seconds): + queue_budget = budget_remaining(self.queue_timeout_seconds) + deadline = time.monotonic() + queue_budget + if not self._connection_mutex.acquire(timeout=queue_budget): raise TimeoutError handle = None try: @@ -98,7 +102,17 @@ def _acquire_connection_slot(self): def _connection_slot(self, timing): try: handle = timing.run("queue", self._acquire_connection_slot) + except RequestBudgetExceeded: + timing.error_category = "request_budget_exhausted" + timing.failed_phase = "queue" + raise except TimeoutError as exc: + try: + budget_remaining(self.queue_timeout_seconds) + except RequestBudgetExceeded: + timing.error_category = "request_budget_exhausted" + timing.failed_phase = "queue" + raise timing.error_category = "journal_queue_timeout" timing.failed_phase = "queue" raise AdapterError(503, "Idempotency journal is busy; retry Add") from exc @@ -121,7 +135,7 @@ def connect(self, operation="probe"): with self.metrics.operation(operation) as timing, self._connection_slot(timing): db = None try: - db = timing.run("connect", sqlite3.connect, self.path, timeout=5) + db = timing.run("connect", sqlite3.connect, self.path, timeout=budget_remaining(5)) timing.run("setup", db.execute, "PRAGMA synchronous = FULL") try: yield TimedConnection(db, timing) diff --git a/agent_memory_leaderboard/src/memoria.py b/agent_memory_leaderboard/src/memoria.py index a024812..4203148 100644 --- a/agent_memory_leaderboard/src/memoria.py +++ b/agent_memory_leaderboard/src/memoria.py @@ -8,6 +8,7 @@ import httpx from pydantic import BaseModel, ConfigDict, Field, ValidationError +from .budget import remaining as budget_remaining from .config import Settings from .errors import AdapterError, VisibilityPending from .models import Nonblank @@ -43,6 +44,7 @@ def _visibility_timeout(self) -> AdapterError: return AdapterError(503, "Memoria write visibility timed out; retry Add", retry_after=delay) def _remaining_visibility(self, deadline: float) -> float: + budget_remaining(float("inf")) remaining = deadline - time.monotonic() if remaining <= 0: raise self._visibility_timeout() @@ -57,12 +59,13 @@ def request( *, deadline: float | None = None, ) -> Any: + budget_remaining(float("inf")) if deadline is not None: self._remaining_visibility(deadline) visibility_limited = False try: with self.scheduler.slot(deadline=deadline): - timeout = self.settings.timeout_seconds + timeout = budget_remaining(self.settings.timeout_seconds) if deadline is not None: remaining = self._remaining_visibility(deadline) visibility_limited = remaining <= timeout @@ -75,9 +78,14 @@ def request( "X-User-Id": owner, }, json=body, - timeout=timeout, + timeout=httpx.Timeout( + timeout, + connect=min(self.settings.connect_timeout_seconds, timeout), + pool=min(self.settings.connect_timeout_seconds, timeout), + ), ) except httpx.TimeoutException as exc: + budget_remaining(float("inf")) if deadline is not None and (visibility_limited or time.monotonic() >= deadline): raise self._visibility_timeout() from exc if isinstance(exc, (httpx.ConnectTimeout, httpx.PoolTimeout)): diff --git a/agent_memory_leaderboard/src/probes.py b/agent_memory_leaderboard/src/probes.py index a3bcc50..6985316 100644 --- a/agent_memory_leaderboard/src/probes.py +++ b/agent_memory_leaderboard/src/probes.py @@ -135,7 +135,12 @@ def observe(event): begin = time.monotonic() try: - with httpx.Client(timeout=args.timeout, trust_env=False) as client: + with httpx.Client( + timeout=httpx.Timeout( + args.timeout, connect=min(30, args.timeout), pool=min(30, args.timeout) + ), + trust_env=False, + ) as client: result = call_adapter( client, adapter_url, @@ -168,7 +173,12 @@ def observe(event): def explain(user, query): if not settings: return None - with httpx.Client(timeout=args.timeout, trust_env=False) as client: + with httpx.Client( + timeout=httpx.Timeout( + args.timeout, connect=min(30, args.timeout), pool=min(30, args.timeout) + ), + trust_env=False, + ) as client: response = client.post( settings.memoria_url + "/v1/memories/search", headers={ diff --git a/agent_memory_leaderboard/src/scheduling.py b/agent_memory_leaderboard/src/scheduling.py index 5c6db78..6c27aaa 100644 --- a/agent_memory_leaderboard/src/scheduling.py +++ b/agent_memory_leaderboard/src/scheduling.py @@ -8,6 +8,7 @@ import anyio +from .budget import remaining as budget_remaining from .errors import AdapterError @@ -34,9 +35,10 @@ def __init__(self, capacity: int, queue_timeout: float): @asynccontextmanager async def slot(self): try: - with anyio.fail_after(self.queue_timeout): + with anyio.fail_after(budget_remaining(self.queue_timeout)): await self.permits.acquire() except TimeoutError as exc: + budget_remaining(self.queue_timeout) depth = self.permits.statistics().tasks_waiting + self.capacity + 1 raise AdapterError( 503, @@ -75,11 +77,12 @@ def _wake_head(self) -> None: @contextmanager def slot(self, *, deadline: float | None = None): if not self.concurrency and not self.interval: + budget_remaining(self.queue_timeout) yield return ticket = Event() deadline = min( - time.monotonic() + self.queue_timeout, + time.monotonic() + budget_remaining(self.queue_timeout), deadline if deadline is not None else math.inf, ) acquired = False @@ -92,6 +95,7 @@ def slot(self, *, deadline: float | None = None): now = time.monotonic() remaining = deadline - now if remaining <= 0: + budget_remaining(self.queue_timeout) depth = len(self.waiters) + self.active raise AdapterError( 503, diff --git a/agent_memory_leaderboard/src/service.py b/agent_memory_leaderboard/src/service.py index d9ba92f..f508ebd 100644 --- a/agent_memory_leaderboard/src/service.py +++ b/agent_memory_leaderboard/src/service.py @@ -2,6 +2,8 @@ import json from itertools import chain, pairwise +from .budget import remaining as budget_remaining +from .budget import request_budget from .config import Settings from .errors import AdapterError from .ingestion import message_parts @@ -28,7 +30,7 @@ def owner(self, user_id: str) -> str: return digest([self.settings.namespace, user_id]) def add(self, request: AddRequest) -> AddResponse: - with self.journal.metrics.add(): + with request_budget(self.settings.request_timeout_seconds), self.journal.metrics.add(): owner = self.owner(request.user_id) key = digest([self.settings.memoria_url.rstrip("/"), owner, request.request_id]) fingerprint = digest(request.model_dump()) @@ -41,6 +43,7 @@ def add(self, request: AddRequest) -> AddResponse: if prepared is None: break self.write(owner, key, position, body, prepared, last=following is None) + budget_remaining(float("inf")) return AddResponse( request_id=request.request_id, user_id=request.user_id, @@ -122,6 +125,12 @@ def write( ) def search(self, request: SearchRequest) -> SearchResponse: + with request_budget(self.settings.request_timeout_seconds): + result = self._search(request) + budget_remaining(float("inf")) + return result + + def _search(self, request: SearchRequest) -> SearchResponse: memories = self.memoria.search( self.owner(request.user_id), text_content(request.query), request.top_k ) diff --git a/agent_memory_leaderboard/src/smoke.py b/agent_memory_leaderboard/src/smoke.py index 9f49daf..5a9459b 100644 --- a/agent_memory_leaderboard/src/smoke.py +++ b/agent_memory_leaderboard/src/smoke.py @@ -80,7 +80,7 @@ def run(adapter_url: str) -> dict: run_id = uuid.uuid4().hex users = [f"aml-smoke:{run_id}:{suffix}" for suffix in ("alice", "bob", "empty")] headers = {"Authorization": f"Bearer {settings.api_key}"} if settings.api_key else {} - with httpx.Client(timeout=settings.timeout_seconds, trust_env=False) as http: + with httpx.Client(timeout=httpx.Timeout(1800, connect=30, pool=30), trust_env=False) as http: def backend(method, path, user, **kwargs): response = http.request( diff --git a/agent_memory_leaderboard/tests/test_request_budget.py b/agent_memory_leaderboard/tests/test_request_budget.py new file mode 100644 index 0000000..c886f5e --- /dev/null +++ b/agent_memory_leaderboard/tests/test_request_budget.py @@ -0,0 +1,174 @@ +"""Synthetic request budget and durable recovery checks.""" + +import sqlite3 +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from types import SimpleNamespace + +import httpx +import pytest +from test_endpoints import FakeMemoria, addition, running + +from src import budget +from src.budget import RequestBudgetExceeded, remaining, request_budget +from src.config import Settings +from src.journal import Journal +from src.memoria import IDEMPOTENT_PATH, Memoria +from src.models import AddRequest +from src.service import MemoryAdapter + + +@pytest.fixture +def setup(tmp_path): + settings = Settings( + memoria_url="http://memoria.test", + memoria_master_key="upstream-secret", + state_dir=tmp_path, + request_timeout_seconds=1, + ) + upstream = FakeMemoria() + return settings, upstream + + +def test_timeout_defaults_and_environment(monkeypatch): + monkeypatch.setenv("MEMORIA_URL", "http://memoria.test") + monkeypatch.setenv("MEMORIA_MASTER_KEY", "synthetic") + settings = Settings.from_env() + assert settings.connect_timeout_seconds == settings.journal_queue_timeout_seconds == 30 + assert settings.visibility_timeout_seconds == 60 + assert settings.request_timeout_seconds == 1700 + monkeypatch.setenv("MEMORIA_CONNECT_TIMEOUT_SECONDS", "45") + monkeypatch.setenv("ADAPTER_REQUEST_TIMEOUT_SECONDS", "1600") + settings = Settings.from_env() + assert settings.connect_timeout_seconds == 45 and settings.request_timeout_seconds == 1600 + + +@pytest.mark.parametrize( + "field,maximum", [("connect_timeout_seconds", 120), ("request_timeout_seconds", 1740)] +) +@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf"), 1800]) +def test_invalid_budget_rejected(setup, field, maximum, value): + with pytest.raises(ValueError): + replace(setup[0], **{field: value}) + + +def test_nested_budget_cannot_extend_parent(monkeypatch): + clock = [100.0] + monkeypatch.setattr(budget, "time", SimpleNamespace(monotonic=lambda: clock[0])) + with request_budget(1): + clock[0] += 0.4 + with request_budget(10): + assert remaining(20) == pytest.approx(0.6) + clock[0] += 0.7 + with pytest.raises(RequestBudgetExceeded): + remaining(20) + assert remaining(20) == 20 + + +def test_budget_spans_chunks_and_retry_preserves_receipts(setup, monkeypatch): + settings, upstream = setup + clock = [100.0] + monkeypatch.setattr(budget, "time", SimpleNamespace(monotonic=lambda: clock[0])) + + def consume_time(request, *_): + if request.url.path == IDEMPOTENT_PATH: + clock[0] += 0.6 + + upstream.hook = consume_time + payload = AddRequest.model_validate( + addition(messages=[{"role": "user", "content": f"Synthetic fact {i}"} for i in range(3)]) + ) + with httpx.Client(transport=httpx.MockTransport(upstream)) as client: + adapter = MemoryAdapter(settings, Memoria(settings, client), Journal(settings.state_dir)) + with pytest.raises(RequestBudgetExceeded): + adapter.add(payload) + assert len(upstream.writes) == 2 # No third write after shared budget expires. + with sqlite3.connect(adapter.journal.path) as db: + assert db.execute("SELECT complete FROM requests").fetchone() == (0,) + upstream.hook = None + adapter.add(payload) + assert len(upstream.memories) == 3 + assert len(upstream.receipts) == 3 + + +def test_journal_wait_consumes_request_budget_without_leaking_lock(setup): + journal = Journal(setup[0].state_dir, queue_timeout_seconds=30) + journal._connection_mutex.acquire() + try: + with request_budget(0.03), pytest.raises(RequestBudgetExceeded): + journal.state("synthetic", 0) + finally: + journal._connection_mutex.release() + assert journal.begin("synthetic", "fingerprint") is False + + +def test_admission_budget_expires_while_original_write_keeps_its_permit(setup): + settings, upstream = setup + settings = replace(settings, max_active_requests=1, request_timeout_seconds=0.05) + entered, release = threading.Event(), threading.Event() + + def block(request, *_): + if request.url.path == IDEMPOTENT_PATH: + entered.set() + assert release.wait(3) + + upstream.hook = block + with running(settings, upstream) as client, ThreadPoolExecutor(max_workers=1) as pool: + original = pool.submit(client.post, "/add", json=addition()) + try: + assert entered.wait(2) + waiting = client.post("/add", json=addition(request_id="other-synthetic-request")) + assert waiting.status_code == 504 + assert not original.done() + assert client.get("/health").status_code == 200 + finally: + release.set() + assert original.result(timeout=3).status_code == 504 + upstream.hook = None + assert client.post("/add", json=addition()).status_code == 200 + assert len(upstream.memories) == 1 + + +def test_connection_and_io_timeouts_are_separate(setup): + settings, upstream = setup + seen = [] + + def inspect_timeout(request, *_): + seen.append(request.extensions["timeout"]) + + upstream.hook = inspect_timeout + with httpx.Client(transport=httpx.MockTransport(upstream)) as client: + Memoria(settings, client).search("a" * 64, "synthetic", 1) + assert seen[0] == {"connect": 30, "pool": 30, "read": 120, "write": 120} + + +def test_expiry_after_final_commit_preserves_completed_replay(setup, monkeypatch): + settings, upstream = setup + clock = [100.0] + monkeypatch.setattr(budget, "time", SimpleNamespace(monotonic=lambda: clock[0])) + journal = Journal(settings.state_dir) + original = sqlite3.connect + commits = 0 + + class SlowFinalCommit(sqlite3.Connection): + def commit(self): + nonlocal commits + super().commit() + commits += 1 + if commits == 3: + clock[0] += 2 + + monkeypatch.setattr( + sqlite3, "connect", lambda *a, **kw: original(*a, factory=SlowFinalCommit, **kw) + ) + with httpx.Client(transport=httpx.MockTransport(upstream)) as client: + adapter = MemoryAdapter(settings, Memoria(settings, client), journal) + request = AddRequest.model_validate(addition()) + with pytest.raises(RequestBudgetExceeded): + adapter.add(request) + with sqlite3.connect(journal.path) as db: + assert db.execute("SELECT complete FROM requests").fetchone() == (1,) + calls = len(upstream.calls) + adapter.add(request) + assert len(upstream.calls) == calls