diff --git a/src/schematic/client.py b/src/schematic/client.py index 63f6957..a6f3309 100644 --- a/src/schematic/client.py +++ b/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 diff --git a/tests/custom/test_client.py b/tests/custom/test_client.py index b02cc7d..7b842a1 100644 --- a/tests/custom/test_client.py +++ b/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()