pin the reservation to the lease try_reserve actually charged

conformance/SPEC.md

@@ -70,7 +70,7 @@ Store-level (exercise the lease store and reservation store directly):
 | `advance_clock` | `ms` | — |
 | `replace_lease` | `lease_id`, `company_id`, `credit_type_id`, `granted_amount`, `expires_at_ms` | `written` (bool) |
 | `drop_lease` | `company_id`, `credit_type_id` | — |
-| `try_reserve` | `company_id`, `credit_type_id`, `credits` | `balance` (post-debit number, or `null`) |
+| `try_reserve` | `company_id`, `credit_type_id`, `credits` | `balance` (post-debit number, or `null`), `lease_id`? (the lease actually charged) |
 | `refund_lease` | `company_id`, `credit_type_id`, `credits`, `pin_lease_id`? | — |
 | `extend_lease` | `company_id`, `credit_type_id`, `granted_total`, `expires_at_ms`?, `pin_lease_id`? | — |
 | `get_lease` | `company_id`, `credit_type_id` | `exists`, `lease_id`?, `granted_amount`?, `local_remaining_credits`? |
@@ -178,7 +178,14 @@ Returns written/kept so the caller can run the redundant-lease release logic (se
   never reach the arithmetic — it slips through every comparison and would poison the balance
   into approving everything).
 - Otherwise debit and return the **post-debit balance** (so the caller can derive the pre-debit
-  figure as `returned + credits` without a racy follow-up read).
+  figure as `returned + credits` without a racy follow-up read) **and the `lease_id` of the lease
+  that was charged**, read atomically with the debit (in-process: under the same per-slot lock;
+  Redis: inside the same script).
+- The debit is **not keyed by lease id** — it charges whichever lease occupies the slot at that
+  moment, which need not be the one the caller's acquire returned (the slot's lease can be
+  replaced in between by expiry + a successor install, by the sweeper, or by a sibling process
+  sharing the backend). Reporting the charged id is what lets the caller pin its reservation to
+  the lease it actually drew on; see [check flow](#check-flow) step 9.
 - Reserving down to exactly 0 is allowed.
 
 ### `refund(company, credit, credits, pin_lease_id?)`
@@ -319,9 +326,14 @@ Then:
    Still refused → failure handling, reason `insufficient_lease_balance`. Store error →
    failure handling, reason `lease_store_error`.
 9. **Record the reservation** (TTL = `reservation_ttl_ms` from now) — *after* the debit, *before*
-   the engine gate. If persisting fails, undo the debit (claim-and-refund; direct refund if
-   nothing persisted; both pinned to the lease) and go to failure handling
-   (`lease_store_error`). If even the undo fails, accept the bounded leak.
+   the engine gate. Pin it to the `lease_id` **`try_reserve` reported**, never to the one the
+   acquire in step 7 returned: those differ whenever the slot's lease was replaced in the window
+   between them (which spans the awaited extend in step 8), and a stale pin sends the settle
+   refund and the sweep refund to a lease that was never charged — `refund`'s pin drops both —
+   while the Track event bills a released lease. If persisting fails, undo the debit
+   (claim-and-refund; direct refund if nothing persisted; both pinned to the charged lease) and
+   go to failure handling (`lease_store_error`). If even the undo fails, accept the bounded
+   leak.
 10. **Engine gate.** Re-run the engine against a company snapshot whose
     `credit_balances[credit_id]` is substituted with the **pre-reservation** local balance
     (post-debit balance returned by `try_reserve` + `credit_cost` — exact as of the debit, no

conformance/vectors/lease-lifecycle.json

@@ -221,6 +221,56 @@
                 }
             ]
         },
+        {
+            "name": "try_reserve_reports_the_lease_it_charged",
+            "description": "try_reserve is not keyed by lease id: it debits whichever lease holds the slot. When the incumbent has expired and a successor was installed, the debit lands on the successor, and try_reserve must report the successor's id \u2014 that is the id the caller pins its reservation to, so the settle and sweep refunds (both pinned) apply and the Track event bills the lease that was actually charged.",
+            "given": {
+                "leases": [
+                    {
+                        "lease_id": "lse_1",
+                        "company_id": "co_1",
+                        "credit_type_id": "ct_1",
+                        "granted_amount": 1000,
+                        "expires_at_ms": 60000
+                    }
+                ]
+            },
+            "operations": [
+                {
+                    "op": "advance_clock",
+                    "ms": 60001
+                },
+                {
+                    "op": "replace_lease",
+                    "lease_id": "lse_2",
+                    "company_id": "co_1",
+                    "credit_type_id": "ct_1",
+                    "granted_amount": 500,
+                    "expires_at_ms": 180000,
+                    "expect": { "written": true }
+                },
+                {
+                    "op": "try_reserve",
+                    "company_id": "co_1",
+                    "credit_type_id": "ct_1",
+                    "credits": 100,
+                    "expect": { "balance": 400, "lease_id": "lse_2" }
+                },
+                {
+                    "op": "refund_lease",
+                    "company_id": "co_1",
+                    "credit_type_id": "ct_1",
+                    "credits": 100,
+                    "pin_lease_id": "lse_2"
+                },
+                {
+                    "op": "get_lease",
+                    "company_id": "co_1",
+                    "credit_type_id": "ct_1",
+                    "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 500 }
+                }
+            ]
+        },
         {
             "name": "refund_clamped_at_granted_amount",
             "description": "A refund can never push the local balance above the granted amount.",

src/schematic/leases/__init__.py

@@ -10,7 +10,7 @@
 
 from .check import CheckDataStream, CreditCheckDeps, check_with_lease
 from .lease_manager import CreditsWireClient, LeaseGrant, LeaseManager, LeaseWireClient
-from .lease_store import InMemoryLeaseStore, LeaseStore, lease_key
+from .lease_store import InMemoryLeaseStore, LeaseStore, ReserveResult, lease_key
 from .redis_lease_store import RedisLeaseStore
 from .redis_reservation_store import RedisReservationStore
 from .reservation_store import InMemoryReservationStore, ReservationStore
@@ -62,6 +62,7 @@
     "ReservationConsumeResult",
     "ReservationRecord",
     "ReservationStore",
+    "ReserveResult",
     "ResolvedLeaseConfig",
     "build_reservation_track_event",
     "check_with_lease",

src/schematic/leases/check.py

@@ -206,18 +206,22 @@ async def failure(reason: str) -> "CheckResult":
         return await failure("lease_acquire_failed")
 
     # try_reserve is the atomic gate: check and debit in one step, returning
-    # the post-debit balance so the pre-debit figure needs no second read.
+    # the post-debit balance so the pre-debit figure needs no second read, plus
+    # the id of the lease it charged. That id, not lease.lease_id, is what the
+    # reservation pins to: the debit is not keyed by lease, so the slot's lease
+    # may have been replaced since the acquire above, over a window that spans
+    # the awaited extend below.
     try:
-        post_reserve_balance = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost)
-        if post_reserve_balance is None:
+        reserve = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost)
+        if reserve is None:
             # Pass the cost as required_credits so a single large request
             # extends even while the ratio sits above the water mark.
             await deps.manager.maybe_extend(resolved_company.id, credit_id, credit_cost, options.timeout)
-            post_reserve_balance = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost)
+            reserve = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost)
     except Exception as err:
         log.error(f"Lease check: reserve against {resolved_company.id}/{credit_id} failed: {err}")
         return await failure("lease_store_error")
-    if post_reserve_balance is None:
+    if reserve is None:
         return await failure("insufficient_lease_balance")
 
     # Record the hold after the debit and before the gate. A crash between the
@@ -227,7 +231,11 @@ async def failure(reason: str) -> "CheckResult":
     resolved_config = deps.manager.resolve_config(credit_id)
     record = ReservationRecord(
         id=str(uuid.uuid4()),
-        lease_id=lease.lease_id,
+        # The lease the debit actually landed on, which may not be the one
+        # acquire_if_needed handed back. Pinning the acquired id instead would
+        # send the settle refund, the sweep refund, and the track event's lease
+        # id to a lease that was never charged.
+        lease_id=reserve.lease_id,
         company_id=resolved_company.id,
         credit_type_id=credit_id,
         event_subtype=event_subtype,
@@ -244,12 +252,14 @@ async def failure(reason: str) -> "CheckResult":
         log.error(f"Lease check: failed to persist reservation {record.id}: {err}")
         # Undo the debit rather than strand it until lease expiry. consume
         # claims whatever slice of the add landed and refunds it; a None says
-        # nothing landed, so refund the debit directly. Both are pinned to this
-        # lease. If the undo itself fails, accept the bounded leak: the slice
-        # comes back at lease expiry, which beats risking a double refund.
+        # nothing landed, so refund the debit directly. Both are pinned to the
+        # lease the debit landed on (the record carries that id, so consume
+        # pins to it too), never to the acquired one. If the undo itself fails,
+        # accept the bounded leak: the slice comes back at lease expiry, which
+        # beats risking a double refund.
         try:
             if await deps.reservations.consume(record.id, 0) is None:
-                await deps.lease_store.refund(resolved_company.id, credit_id, credit_cost, lease.lease_id)
+                await deps.lease_store.refund(resolved_company.id, credit_id, credit_cost, reserve.lease_id)
         except Exception as undo_err:
             log.warning(
                 f"Lease check: could not undo the local debit for {record.id} ({undo_err}); "
@@ -262,7 +272,7 @@ async def failure(reason: str) -> "CheckResult":
     # returned plus what it debited, exact as of the debit), and credit_cost
     # tells the engine what this call costs, so it evaluates the same
     # arithmetic try_reserve just enforced, plus every non-credit rule.
-    pre_reservation = post_reserve_balance + credit_cost
+    pre_reservation = reserve.balance + credit_cost
     substituted = _substitute_credit_balance(resolved_company, credit_id, pre_reservation)
     try:
         result = datastream.evaluate_flag(

src/schematic/leases/lease_store.py

@@ -12,6 +12,7 @@
 import math
 import time
 from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from typing import AsyncIterator, Dict, List, Optional, Tuple
 
 from .types import Clock, LeaseState
@@ -21,6 +22,15 @@ def lease_key(company_id: str, credit_type_id: str) -> str:
     return f"{company_id}:{credit_type_id}"
 
 
+@dataclass(frozen=True)
+class ReserveResult:
+    """What a successful ``try_reserve`` reports: the post-debit balance, and
+    the lease the credits actually came out of."""
+
+    balance: float
+    lease_id: str
+
+
 class LeaseStore(abc.ABC):
     """Backing store for lease slots, shared by the in-memory and Redis backends."""
 
@@ -52,13 +62,22 @@ async def replace(
         """
 
     @abc.abstractmethod
-    async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]:
-        """Atomically check and debit, returning the post-debit balance.
+    async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[ReserveResult]:
+        """Atomically check and debit, reporting the post-debit balance and the
+        lease the debit landed on.
 
         ``None`` when there is no lease, it has expired, the balance is short,
         or ``credits`` is not a finite non-negative number. Returning the
         balance (rather than a bool) lets the caller derive the pre-debit
         figure as ``returned + credits`` without a racy follow-up read.
+
+        The debit is NOT keyed by lease id: it charges whichever lease occupies
+        the slot at that moment, which need not be the one the caller's acquire
+        handed back, since the slot's lease can be replaced in between by the
+        sweeper or by a sibling pod on a shared backend. The caller must
+        therefore pin its reservation to the returned ``lease_id``, never to
+        the acquired one: the settle refund, the sweep refund, and the track
+        event's lease id all have to name the lease that was actually charged.
         """
 
     @abc.abstractmethod
@@ -181,7 +200,7 @@ async def replace(
             )
             return True
 
-    async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]:
+    async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[ReserveResult]:
         # NaN passes every comparison below, and a NaN balance would approve
         # every later reserve, so it never reaches the arithmetic.
         if not is_finite_non_negative(credits):
@@ -196,7 +215,10 @@ async def try_reserve(self, company_id: str, credit_type_id: str, credits: float
             if entry.local_remaining_credits < credits:
                 return None
             entry.local_remaining_credits -= credits
-            return entry.local_remaining_credits
+            # The lease id is read under the SAME lock as the debit: the caller
+            # pins its reservation to it, and a read taken after the lock could
+            # name a lease that replaced this one in between.
+            return ReserveResult(balance=entry.local_remaining_credits, lease_id=entry.lease_id)
 
     async def refund(
         self,

src/schematic/leases/redis_lease_store.py

@@ -10,7 +10,7 @@
 import time
 from typing import Any, Dict, List, Optional
 
-from .lease_store import LeaseStore, is_finite_non_negative, lease_key
+from .lease_store import LeaseStore, ReserveResult, is_finite_non_negative, lease_key
 from .types import DEFAULT_LEASE_DURATION, Clock, LeaseState
 
 DEFAULT_KEY_PREFIX = "schematic:"
@@ -86,25 +86,35 @@
 """
 )
 
-# Atomic check-and-decrement on `localRemainingCredits`. Returns the post-debit
-# balance as a string (a Lua number reply truncates to integer, which would
-# corrupt fractional credit costs); nil if there is no lease, the lease has
-# expired, or there is insufficient remaining. The expiry guard compares
-# against the Redis server clock, so a reserve against an expired-but-not-yet-
-# evicted row during the TTL grace window is rejected.
+# Atomic check-and-decrement on `localRemainingCredits`. Returns
+# `{post-debit balance, charged leaseId}`, the balance as a string (a Lua
+# number reply truncates to integer, which would corrupt fractional credit
+# costs); nil if there is no lease, the lease has expired, or there is
+# insufficient remaining. Reading the leaseId inside the same script is what
+# lets the caller pin its reservation to the lease the debit actually landed
+# on: the debit is not keyed by lease id, and a sibling pod can replace the
+# slot's lease at any point before it. The expiry guard compares against the
+# Redis server clock, so a reserve against an expired-but-not-yet-evicted row
+# during the TTL grace window is rejected.
+#
+# The reply shape is safe for a mixed fleet: every pod EVALs its own copy of
+# this script text and decodes its own reply, and the key layout and hash
+# fields are untouched, so old and new pods keep sharing one lease hash.
 TRY_RESERVE_SCRIPT = (
     LEASE_NOW_MS
     + """
 local raw = redis.call('HGET', KEYS[1], 'localRemainingCredits')
 if not raw then return false end
+local lease_id = redis.call('HGET', KEYS[1], 'leaseId')
+if not lease_id then return false end
 local expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0')
 if expiry <= now then return false end
 local remaining = tonumber(raw)
 local requested = tonumber(ARGV[1])
 if remaining < requested then return false end
 local new_remaining = remaining - requested
 redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_remaining))
-return tostring(new_remaining)
+return { tostring(new_remaining), lease_id }
 """
 )
 
@@ -257,7 +267,7 @@ async def replace(
         )
         return _to_number(result) == 1
 
-    async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]:
+    async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[ReserveResult]:
         # Reject non-finite/negative debits before they reach the script: the
         # string form of NaN parses back to a Lua nan, slips through the `<`
         # comparison, and would poison the SHARED balance for every pod.
@@ -268,9 +278,12 @@ async def try_reserve(self, company_id: str, credit_type_id: str, credits: float
             [self.hash_key(company_id, credit_type_id)],
             [format_amount(credits)],
         )
-        if result is None or result is False:
+        # A nil reply (could not reserve) surfaces as None; success is a
+        # two-element multi-bulk of the post-debit balance and the charged
+        # lease id, both as strings.
+        if not result:
             return None
-        return float(to_str(result))
+        return ReserveResult(balance=float(to_str(result[0])), lease_id=to_str(result[1]))
 
     async def refund(
         self,

tests/conformance/test_vectors.py

@@ -165,9 +165,13 @@ async def _op_drop_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any])
 
 
 async def _op_try_reserve(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None:
-    balance = await h.leases.try_reserve(op["company_id"], op["credit_type_id"], op["credits"])
+    reserve = await h.leases.try_reserve(op["company_id"], op["credit_type_id"], op["credits"])
     if "balance" in expect:
-        _assert_number(balance, expect["balance"])
+        _assert_number(reserve.balance if reserve is not None else None, expect["balance"])
+    # Optional: the lease the debit actually landed on, which the caller must
+    # pin its reservation to.
+    if "lease_id" in expect:
+        assert reserve is not None and reserve.lease_id == expect["lease_id"]
 
 
 async def _op_refund_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None:

tests/lease_support.py

@@ -14,7 +14,7 @@
 import fakeredis.aioredis
 
 from schematic.leases import LeaseGrant, LeaseState, ReservationRecord
-from schematic.leases.lease_store import LeaseStore
+from schematic.leases.lease_store import LeaseStore, ReserveResult
 from schematic.types import (
     RulesengineCheckFlagResult,
     RulesengineCompany,
@@ -104,7 +104,9 @@ async def replace(
             expires_at=expires_at,
         )
 
-    async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]:
+    async def try_reserve(
+        self, company_id: str, credit_type_id: str, credits: float
+    ) -> Optional[ReserveResult]:
         return await self._target.try_reserve(company_id, credit_type_id, credits)
 
     async def refund(
@@ -148,6 +150,9 @@ def __init__(self) -> None:
         # Runs while an acquire is in flight, for emulating a sibling pod
         # winning the race.
         self.during_acquire: Optional[Any] = None
+        # The same seam on the extend, for emulating the slot's lease being
+        # replaced while a check waits on the extend wire call.
+        self.during_extend: Optional[Any] = None
 
     async def acquire(
         self,
@@ -195,6 +200,10 @@ async def extend(
                 "timeout": timeout,
             }
         )
+        during = self.during_extend
+        if during is not None:
+            self.during_extend = None
+            await during()
         scripted = self.extend_responses.pop(0) if self.extend_responses else None
         lease = _scripted_lease(scripted, "unscripted extend wire call")
         return LeaseGrant(