fix(gateway): harden streaming settlement and admission

ee/src/shim_enterprise/application.py

@@ -135,6 +135,7 @@ async def _lifespan(application: FastAPI) -> AsyncIterator[None]:
     try:
         yield
     finally:
+        await application.state.gateway_service.kernel.postprocessor.drain()
         await http_client.aclose()
         await cache.close()
         await engine.dispose()

ee/src/shim_enterprise/gateway/pipeline/quota_reservation.py

@@ -97,7 +97,7 @@ async def quota(
             select(TierDefinition)
             .where(TierDefinition.slug == api_key.tier)
             .execution_options(populate_existing=True)
-            .with_for_update()
+            .with_for_update(read=True)
         )
         tier = (await session.execute(tier_statement)).scalar_one_or_none()
         if tier is None:

ee/tests/architecture/test_manual_test_dashboard.py

@@ -86,7 +86,7 @@ async def test_api_lifespan_does_not_start_continuous_reconciliation(
     monkeypatch,
 ) -> None:
     cache = SimpleNamespace(close=AsyncMock())
-    kernel = SimpleNamespace()
+    kernel = SimpleNamespace(postprocessor=SimpleNamespace(drain=AsyncMock()))
     create_gateway_kernel = Mock(return_value=kernel)
     engine = SimpleNamespace(dispose=AsyncMock())
     connect_cache = AsyncMock()
@@ -110,6 +110,7 @@ async def test_api_lifespan_does_not_start_continuous_reconciliation(
     assert create_gateway_kernel.call_args.args[1].is_closed
     connect_cache.assert_awaited_once_with(cache)
     create_task.assert_not_called()
+    kernel.postprocessor.drain.assert_awaited_once()
     cache.close.assert_awaited_once()
     engine.dispose.assert_awaited_once()
     shutdown_tracing.assert_called_once()

ee/tests/gateway/kernel/test_accounting_coordinator.py

@@ -1506,3 +1506,148 @@ async def test_audit_intent_outbox_reference_is_tenant_scoped(db) -> None:
                     },
                 )
         assert "fk_audit_intent_org_outbox_event" in str(reference_error.value.orig)
+
+
+@pytest.mark.asyncio
+async def test_quota_policy_uses_shared_tier_lock_and_exclusive_key_lock():
+    from sqlalchemy.dialects import postgresql
+    from shim_enterprise.gateway.pipeline.quota_reservation import (
+        AccountingPolicyLoader,
+    )
+
+    session = SimpleNamespace(
+        execute=AsyncMock(
+            side_effect=[
+                SimpleNamespace(
+                    scalar_one_or_none=lambda: SimpleNamespace(tier="free")
+                ),
+                SimpleNamespace(
+                    scalar_one_or_none=lambda: SimpleNamespace(
+                        slug="free",
+                        daily_request_limit=10,
+                        monthly_request_limit=100,
+                        monthly_token_limit=1000,
+                    )
+                ),
+            ]
+        )
+    )
+    policy = await AccountingPolicyLoader().quota(session, _prepared())
+    statements = [
+        str(call.args[0].compile(dialect=postgresql.dialect()))
+        for call in session.execute.await_args_list
+    ]
+    assert statements[0].endswith("FOR UPDATE")
+    assert statements[1].endswith("FOR SHARE")
+    assert policy.daily_request_limit == 10
+
+
+@pytest.mark.asyncio
+async def test_shared_tier_allows_independent_reservations_but_fences_edits(
+    async_engine,
+):
+    from sqlalchemy import update
+    from shim_enterprise.billing.ledger import QuotaLimitExceeded
+    from shim_enterprise.gateway.pipeline.quota_reservation import (
+        AccountingPolicyLoader,
+    )
+    from shim_enterprise.tenants.models import TierDefinition
+
+    factory = async_sessionmaker(async_engine, expire_on_commit=False)
+    slug = f"lock-test-{uuid4().hex}"
+    async with factory.begin() as setup:
+        setup.add(
+            TierDefinition(
+                slug=slug,
+                name="Lock test",
+                rate_limit_rpm=60,
+                rate_limit_tpm=1000,
+                daily_request_limit=1,
+                monthly_request_limit=1,
+                monthly_token_limit=1000,
+            )
+        )
+        await setup.flush()
+        tenants = [await _create_tenant(setup, label) for label in ("lock-a", "lock-b")]
+        await setup.execute(
+            update(ApiKey)
+            .where(ApiKey.id.in_([key for _, _, key in tenants]))
+            .values(tier=slug)
+        )
+
+    async def reserve(session, tenant):
+        prepared = _prepared()
+        prepared.tenant_id, _, prepared.api_key_id = tenant
+        policy = await AccountingPolicyLoader().quota(session, prepared)
+        now = datetime.now(timezone.utc)
+        await DurableAccountingRepository().reserve_quota(
+            session,
+            QuotaReservationCommand(
+                tenant_id=prepared.tenant_id,
+                api_key_id=prepared.api_key_id,
+                request_id=prepared.request_id,
+                requested_model=prepared.model,
+                source_endpoint="chat.completions",
+                started_at=now,
+                reconciliation_due_at=now + timedelta(minutes=2),
+                estimated_input_tokens=1,
+                maximum_output_tokens=1,
+                policy=policy,
+            ),
+        )
+        return policy
+
+    try:
+        async with factory() as first, factory() as second, factory() as contender:
+            await reserve(first, tenants[0])
+            await second.execute(text("SET LOCAL lock_timeout = '500ms'"))
+            # This must succeed while the first tenant still holds its tier lock.
+            snapshot = await reserve(second, tenants[1])
+            assert snapshot.daily_request_limit == 1
+
+            await contender.execute(text("SET LOCAL lock_timeout = '100ms'"))
+            with pytest.raises(DBAPIError) as same_key:
+                await reserve(contender, tenants[0])
+            assert same_key.value.orig.sqlstate == "55P03"
+            await contender.rollback()
+            await first.commit()
+            with pytest.raises(QuotaLimitExceeded):
+                await reserve(contender, tenants[0])
+            await contender.rollback()
+
+            tier_update = (
+                update(TierDefinition)
+                .where(TierDefinition.slug == slug)
+                .values(daily_request_limit=2, monthly_request_limit=2)
+            )
+            await contender.execute(text("SET LOCAL lock_timeout = '100ms'"))
+            with pytest.raises(DBAPIError) as policy_edit:
+                await contender.execute(tier_update)
+            assert policy_edit.value.orig.sqlstate == "55P03"
+            await contender.rollback()
+            await second.rollback()
+            await contender.execute(tier_update)
+            await contender.commit()
+            revised = await reserve(contender, tenants[1])
+            assert (revised.daily_request_limit, revised.monthly_request_limit) == (
+                2,
+                2,
+            )
+            assert revised.version != snapshot.version
+    finally:
+        async with factory.begin() as cleanup:
+            ids = [organization for organization, _, _ in tenants]
+            for model in (
+                UsageLedger,
+                RequestLifecycle,
+                QuotaPeriodUsage,
+                ApiKey,
+                User,
+            ):
+                await cleanup.execute(
+                    delete(model).where(model.organization_id.in_(ids))
+                )
+            await cleanup.execute(delete(Organization).where(Organization.id.in_(ids)))
+            await cleanup.execute(
+                delete(TierDefinition).where(TierDefinition.slug == slug)
+            )

src/shim/application.py

@@ -121,6 +121,8 @@ async def lifespan(application: FastAPI) -> AsyncIterator[None]:
         try:
             yield
         finally:
+            await application.state.gateway_service.kernel.postprocessor.drain()
+            await usage.aclose()
             if owns_http_client:
                 await client.aclose()
 

src/shim/gateway/admission.py

@@ -6,6 +6,8 @@
 from collections.abc import Callable, Hashable
 from dataclasses import dataclass
 import hashlib
+import heapq
+from itertools import count
 import time
 from typing import Literal, Protocol
 
@@ -39,6 +41,8 @@ def __init__(
         self.max_entries = max_entries
         self.clock = clock
         self.windows: OrderedDict[Hashable, tuple[int, float]] = OrderedDict()
+        self._expirations: list[tuple[float, int, Hashable]] = []
+        self._sequence = count()
 
     def increment(self, key: Hashable, *, amount: int, window_seconds: int) -> int:
         now = self.clock()
@@ -52,15 +56,25 @@ def increment(self, key: Hashable, *, amount: int, window_seconds: int) -> int:
             while len(self.windows) >= self.max_entries:
                 self.windows.popitem(last=False)
             count = amount
-            self.windows[key] = (count, now + window_seconds)
+            expires_at = now + window_seconds
+            self.windows[key] = (count, expires_at)
+            heapq.heappush(self._expirations, (expires_at, next(self._sequence), key))
+            if len(self._expirations) > 2 * self.max_entries:
+                self._expirations = [
+                    (expiry, next(self._sequence), item)
+                    for item, (_, expiry) in self.windows.items()
+                ]
+                heapq.heapify(self._expirations)
             return count
         count = current[0] + amount
         self.windows[key] = (count, current[1])
         return count
 
     def _discard_expired(self, now: float) -> None:
-        for key, (_, expires_at) in tuple(self.windows.items()):
-            if expires_at <= now:
+        while self._expirations and self._expirations[0][0] <= now:
+            expires_at, _, key = heapq.heappop(self._expirations)
+            current = self.windows.get(key)
+            if current is not None and current[1] == expires_at:
                 del self.windows[key]
 
 

src/shim/gateway/pipeline/admission.py

@@ -131,8 +131,7 @@ async def run(self, value: PreparedInference) -> PreparedInference:
                 },
             )
         input_tokens = _estimate_input_tokens(payload)
-        candidate_count = _candidate_count(payload)
-        output_tokens = per_candidate_output_tokens * candidate_count
+        output_tokens = per_candidate_output_tokens * candidate_count(value)
         tier = value.context.tier_policy
         key_hash = value.policy.rate_limit_key_hash
         if tier.rate_limit_rpm is not None and not await self.rate_limiter.allow(
@@ -206,13 +205,16 @@ def _estimate_input_tokens(payload: Mapping[str, object]) -> int:
     return max(1, len(serialized.encode("utf-8", errors="backslashreplace")))
 
 
-def _candidate_count(payload: Mapping[str, object]) -> int:
-    generation_config = payload.get("generationConfig")
-    candidate = (
-        generation_config.get("candidateCount")
-        if isinstance(generation_config, Mapping)
-        else payload.get("n", 1)
-    )
+def candidate_count(prepared: PreparedInference) -> int:
+    if prepared.provider == "google":
+        config = prepared.payload.get("generationConfig")
+        candidate = (
+            config.get("candidateCount", 1) if isinstance(config, Mapping) else 1
+        )
+    elif prepared.provider == "openai" and prepared.protocol == "chat":
+        candidate = prepared.payload.get("n", 1)
+    else:
+        candidate = 1
     count = (
         candidate
         if isinstance(candidate, int) and not isinstance(candidate, bool)

src/shim/gateway/pipeline/anthropic_execution.py

@@ -147,7 +147,8 @@ async def close_stream() -> None:
                     return
                 state["closed"] = True
                 try:
-                    await result.close()
+                    async with asyncio.timeout(5):
+                        await result.close()
                 except Exception:
                     pass
                 finally:
@@ -334,6 +335,7 @@ def _error_event(error: ProviderCallError) -> bytes:
         "type": "error",
         "error": {
             "type": "api_error",
+            "code": error.error_code,
             "message": message,
         },
     }

src/shim/gateway/pipeline/google_execution.py

@@ -19,7 +19,7 @@
     ProviderStream,
 )
 from shim.gateway.streaming.sse import encode_data
-from shim.privacy.deanonymizer import _split_placeholder_prefix
+from shim.privacy.deanonymizer import restore_fragment
 from shim.privacy.pii_scrubber import PIIScrubberService
 from shim.secrets.credentials import ProviderCredentialResolver
 
@@ -206,7 +206,7 @@ async def _stream(
         except (asyncio.CancelledError, GeneratorExit):
             raise
         except Exception as exc:
-            if state["recorded"]:
+            if state["recorded"] and not isinstance(exc, ValueError):
                 return
             await self._record_error(exc)
             state["recorded"] = True
@@ -289,11 +289,9 @@ def _restore_value(
         return value
 
     def _restore_fragment(self, key: tuple[object, ...], fragment: str) -> str:
-        text = self._buffers.pop(key, "") + fragment
-        ready, carry = _split_placeholder_prefix(text, self._verification_map)
-        if carry:
-            self._buffers[key] = carry
-        return self._scrubber.deanonymize(ready, self._verification_map)
+        return restore_fragment(
+            self._buffers, key, fragment, self._verification_map, self._scrubber
+        )
 
     def _flush_candidate(
         self,
@@ -481,7 +479,8 @@ def _stream_error(error: ProviderCallError) -> bytes:
 
 async def _close_client(client: genai.Client) -> None:
     try:
-        await client.aio.aclose()
+        async with asyncio.timeout(5):
+            await client.aio.aclose()
     except Exception:
         pass
     try:
@@ -492,6 +491,7 @@ async def _close_client(client: genai.Client) -> None:
 
 async def _close_stream(stream) -> None:
     try:
-        await stream.aclose()
+        async with asyncio.timeout(5):
+            await stream.aclose()
     except Exception:
         pass

src/shim/gateway/pipeline/openai_execution.py

@@ -131,7 +131,8 @@ async def close_stream() -> None:
                     return
                 state["closed"] = True
                 try:
-                    await result.close()
+                    async with asyncio.timeout(5):
+                        await result.close()
                 except Exception:
                     pass
                 finally:
@@ -326,7 +327,7 @@ async def _chat_stream(
         except (asyncio.CancelledError, GeneratorExit):
             raise
         except Exception as exc:
-            if state["recorded"]:
+            if state["recorded"] and not isinstance(exc, ValueError):
                 yield b"data: [DONE]\n\n"
                 return
             await self._record_error(exc)