send an idempotency key on check-and-reserve and retry again

src/schematic/client.py

@@ -4,6 +4,7 @@
 import logging
 import math
 import time
+import uuid
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
@@ -398,19 +399,20 @@ def _resolve_reservation_ttl(logger: logging.Logger, credit_leases: Optional[Cre
 
 
 def _reservation_request_kwargs(options: CheckOptions) -> Dict[str, Any]:
-    """Preflight body and request options for a check-and-reserve call, with
-    the preflight omitted when the caller set nothing.
-
-    Retries are always off. The default policy re-sends on 408, 429 and 5xx,
-    and this request carries no idempotency key, so a 502 arriving after the
-    server committed the hold would take a second one against the same
-    balance. A caller that wants the call retried can retry the check.
+    """Preflight body, idempotency key and request options for a
+    check-and-reserve call, with the preflight omitted when the caller set
+    nothing.
+
+    One key per check, minted before the call so that every attempt the retry
+    policy makes carries the same one: the server answers the repeat with the
+    hold the first attempt took, so a 502 arriving after it committed no
+    longer costs a second hold.
     """
-    kwargs: Dict[str, Any] = {}
+    kwargs: Dict[str, Any] = {"idempotency_key": str(uuid.uuid4())}
     preflight = _build_preflight(_check_options_to_flag_options(options))
     if preflight is not None:
         kwargs["preflight"] = preflight
-    request_options: RequestOptions = {"max_retries": 0}
+    request_options: RequestOptions = {}
     if options.timeout is not None:
         request_options["timeout"] = options.timeout
     kwargs["request_options"] = request_options

tests/custom/test_client.py

@@ -1,12 +1,13 @@
 import asyncio
 import datetime as dt
+import json
 import time
 import unittest
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from httpx import AsyncClient, Client
+from httpx import AsyncClient, Client, MockTransport, Response
 from lease_support import ScriptedDataStream, ScriptedEngine, make_fake_redis
 
 from schematic.cache import LocalCache, RedisCache
@@ -31,6 +32,7 @@
     TrackWithReservationOptions,
     _is_valid_quantity,
 )
+from schematic.core import http_client as core_http_client
 from schematic.core.api_error import ApiError as CoreApiError
 from schematic.errors import PaymentRequiredError
 from schematic.leases import LeaseConfigOverride
@@ -1748,7 +1750,8 @@ def test_returns_a_reservation_handle_built_from_the_response(self):
         ttl = dt.timedelta(seconds=TTL_SECONDS)
         self.assertGreaterEqual(kwargs["expires_at"], before + ttl)
         self.assertLessEqual(kwargs["expires_at"], after + ttl)
-        self.assertEqual(kwargs["request_options"], {"max_retries": 0})
+        self.assertEqual(kwargs["request_options"], {})
+        self.assertTrue(kwargs["idempotency_key"])
 
         # The server logs the flag check for check-and-reserve itself.
         mock_push.assert_not_called()
@@ -1761,14 +1764,21 @@ def test_sends_the_generic_usage_preflight_without_an_event_subtype(self):
     def test_forwards_the_per_check_timeout(self):
         self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50, timeout=2.5))
         kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs
-        self.assertEqual(kwargs["request_options"], {"max_retries": 0, "timeout": 2.5})
+        self.assertEqual(kwargs["request_options"], {"timeout": 2.5})
 
-    def test_never_retries_check_and_reserve(self):
-        # The call has no idempotency key, so a retried 5xx that the server
-        # already committed would take a second hold.
+    def test_leaves_the_default_retry_policy_in_place(self):
+        # The idempotency key is what makes a retried 5xx safe, so the call no
+        # longer opts out of retries.
         self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
         kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs
-        self.assertEqual(kwargs["request_options"]["max_retries"], 0)
+        self.assertNotIn("max_retries", kwargs["request_options"])
+
+    def test_mints_a_fresh_idempotency_key_per_check(self):
+        self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
+        first = self.schematic.features.check_and_reserve_flag.call_args.kwargs["idempotency_key"]
+        self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
+        second = self.schematic.features.check_and_reserve_flag.call_args.kwargs["idempotency_key"]
+        self.assertNotEqual(first, second)
 
     def test_a_fractional_usage_sizes_the_hold_and_rounds_the_preflight_up(self):
         self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=0.5))
@@ -2208,6 +2218,8 @@ async def test_returns_a_reservation_handle_built_from_the_response(self):
         )
         ttl = dt.timedelta(seconds=TTL_SECONDS)
         assert before + ttl <= kwargs["expires_at"] <= after + ttl
+        assert kwargs["request_options"] == {}
+        assert kwargs["idempotency_key"]
 
         # The server logs the flag check for check-and-reserve itself.
         mock_push.assert_not_called()
@@ -2220,14 +2232,21 @@ async def test_sends_the_generic_usage_preflight_without_an_event_subtype(self):
     async def test_forwards_the_per_check_timeout(self):
         await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50, timeout=2.5))
         kwargs = self.client.features.check_and_reserve_flag.call_args.kwargs
-        assert kwargs["request_options"] == {"max_retries": 0, "timeout": 2.5}
+        assert kwargs["request_options"] == {"timeout": 2.5}
 
-    async def test_never_retries_check_and_reserve(self):
-        # The call has no idempotency key, so a retried 5xx that the server
-        # already committed would take a second hold.
+    async def test_leaves_the_default_retry_policy_in_place(self):
+        # The idempotency key is what makes a retried 5xx safe, so the call no
+        # longer opts out of retries.
         await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
         kwargs = self.client.features.check_and_reserve_flag.call_args.kwargs
-        assert kwargs["request_options"] == {"max_retries": 0}
+        assert "max_retries" not in kwargs["request_options"]
+
+    async def test_mints_a_fresh_idempotency_key_per_check(self):
+        await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
+        first = self.client.features.check_and_reserve_flag.call_args.kwargs["idempotency_key"]
+        await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
+        second = self.client.features.check_and_reserve_flag.call_args.kwargs["idempotency_key"]
+        assert first != second
 
     async def test_a_fractional_usage_sizes_the_hold_and_rounds_the_preflight_up(self):
         await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=0.5))
@@ -3014,5 +3033,102 @@ def test_auto_warns_about_the_client_only_options_too(self):
             client.event_buffer.stop()
 
 
+class _ReserveTransport(MockTransport):
+    """Replays a queue of status codes in order, recording each request body."""
+
+    def __init__(self, statuses):
+        self.bodies = []
+        self._statuses = list(statuses)
+        super().__init__(self._handle)
+
+    def _handle(self, request):
+        self.bodies.append(json.loads(request.content) if request.content else {})
+        status = self._statuses.pop(0) if self._statuses else 200
+        if status >= 400:
+            return Response(status, json={"error": "upstream is unhappy"})
+        return Response(status, json=_reserve_body())
+
+
+def _reserve_body():
+    expires_at = (dt.datetime.now(dt.timezone.utc) + dt.timedelta(seconds=TTL_SECONDS)).isoformat()
+    return {
+        "data": {
+            "flag": "inference",
+            "flag_id": "flag_1",
+            "value": True,
+            "reason": "matched",
+            "company_id": "co_1",
+            "entitlement": {"feature_id": "feat", "feature_key": "inference", "value_type": "credit"},
+            "reservation": {
+                "id": "rsv_1",
+                "company_id": "co_1",
+                "credit_type_id": "bilcr_inference",
+                "consumption_rate": 10.0,
+                "credits_reserved": 500.0,
+                "quantity_reserved": 50.0,
+                "event_subtype": "inference_tokens",
+                "expires_at": expires_at,
+            },
+        },
+        "params": {},
+    }
+
+
+class TestServerReservationRetries(unittest.TestCase):
+    """check() in server mode across the SDK's own retry policy.
+
+    These drive the generated features client over a mock transport rather
+    than a stubbed method, because what they pin is the retry loop itself:
+    which body each attempt carries, and what the check makes of the attempt
+    that finally succeeds.
+    """
+
+    def setUp(self):
+        # The retry policy sleeps a second before its first retry, which no
+        # test needs to sit through.
+        self._delay = patch.object(core_http_client, "INITIAL_RETRY_DELAY_SECONDS", 0.001)
+        self._delay.start()
+        self.addCleanup(self._delay.stop)
+
+    def _client(self, transport) -> Schematic:
+        client = Schematic(
+            "api_key",
+            SchematicConfig(
+                event_buffer_period=1,
+                logger=MagicMock(),
+                httpx_client=Client(transport=transport),
+                credit_leases=CreditLeaseConfig(mode="server", default_reservation_ttl=TTL_SECONDS),
+            ),
+        )
+        self.addCleanup(client.event_buffer.stop)
+        client.flag_check_cache_providers = []
+        return client
+
+    def test_a_retried_check_repeats_its_key_and_holds_once(self):
+        transport = _ReserveTransport([502, 200])
+        client = self._client(transport)
+
+        result = client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
+
+        self.assertEqual(len(transport.bodies), 2)
+        keys = [body["idempotency_key"] for body in transport.bodies]
+        self.assertEqual(keys[0], keys[1])
+        self.assertTrue(result.allowed)
+        assert result.reservation is not None
+        self.assertEqual(result.reservation.id, "rsv_1")
+
+    def test_the_next_check_carries_a_different_key(self):
+        transport = _ReserveTransport([502, 200, 200])
+        client = self._client(transport)
+
+        client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
+        client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50))
+
+        keys = [body["idempotency_key"] for body in transport.bodies]
+        self.assertEqual(len(keys), 3)
+        self.assertEqual(keys[0], keys[1])
+        self.assertNotEqual(keys[1], keys[2])
+
+
 if __name__ == "__main__":
     unittest.main()