From a9d8b5f6648f5963b433ee337f06b9dedbdb02f3 Mon Sep 17 00:00:00 2001 From: foxjl85-dev <272123326+foxjl85-dev@users.noreply.github.com> Date: Mon, 29 Jun 2026 08:07:31 -0500 Subject: [PATCH 1/3] docs(proxy): add duplicate burst guard cookbook example --- .../duplicate_burst_guard.py | 148 ++++++++++++++++++ .../test_duplicate_burst_guard_example.py | 80 ++++++++++ 2 files changed, 228 insertions(+) create mode 100644 cookbook/litellm_proxy_server/duplicate_burst_guard.py create mode 100644 tests/test_litellm/test_duplicate_burst_guard_example.py diff --git a/cookbook/litellm_proxy_server/duplicate_burst_guard.py b/cookbook/litellm_proxy_server/duplicate_burst_guard.py new file mode 100644 index 00000000000..c7d13b26ebb --- /dev/null +++ b/cookbook/litellm_proxy_server/duplicate_burst_guard.py @@ -0,0 +1,148 @@ +""" +Example LiteLLM Proxy callback that blocks short duplicate request bursts. + +Register this handler in proxy_config.yaml: + +litellm_settings: + callbacks: + - duplicate_burst_guard.duplicate_burst_guard + +This is an in-memory, single-process example. Use a shared store if your proxy +runs multiple workers or needs duplicate detection across instances. +""" + +import asyncio +import hashlib +import time +from collections import defaultdict, deque +from collections.abc import Mapping +from typing import Optional, cast + +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral + + +class DuplicateBurstGuard(CustomLogger): + def __init__(self, max_calls: int = 2, window_seconds: float = 10.0) -> None: + self.max_calls = max_calls + self.window_seconds = window_seconds + self._calls: defaultdict[str, deque[float]] = defaultdict( + deque + ) # mutable-ok: sliding window state; use Redis ZSET for multi-worker + self._last_gc = 0.0 # mutable-ok: local GC watermark for in-process state + self._lock = asyncio.Lock() + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + fingerprint = self._fingerprint( + data=data, user_api_key_dict=user_api_key_dict, call_type=call_type + ) + count = await self._record_and_count(fingerprint) + if count > self.max_calls: + raise HTTPException( + status_code=429, + detail={ + "error": "Duplicate request burst detected", + "fingerprint": fingerprint, + }, + ) + return data + + async def _record_and_count(self, fingerprint: str) -> int: + now = time.monotonic() + async with self._lock: + if now - self._last_gc > self.window_seconds: + self._prune(now) + self._last_gc = now + + timestamps = self._calls[fingerprint] + timestamps.append(now) + self._prune_key(timestamps, now) + return len(timestamps) + + def _prune(self, now: float) -> None: + for fingerprint, timestamps in list(self._calls.items()): + self._prune_key(timestamps, now) + if not timestamps: + self._calls.pop(fingerprint, None) + + def _prune_key(self, timestamps: deque[float], now: float) -> None: + while timestamps and now - timestamps[0] > self.window_seconds: + timestamps.popleft() + + def _fingerprint( + self, + data: dict[str, object], + user_api_key_dict: Optional[UserAPIKeyAuth], + call_type: CallTypesLiteral, + ) -> str: + messages_value = data.get("messages") + system_prompt = "" + last_user_prompt = "" + prompt = "" + + if isinstance(messages_value, list): + messages = cast(list[object], messages_value) + for message in messages: + if not isinstance(message, Mapping): + continue + message_data = cast(Mapping[str, object], message) + if message_data.get("role") == "system" and not system_prompt: + system_prompt = self._content_text(message_data.get("content")) + if message_data.get("role") == "user": + last_user_prompt = self._content_text(message_data.get("content")) + else: + prompt = self._content_text(data.get("prompt")) + + metadata_value = data.get("metadata") + metadata: Mapping[str, object] + if isinstance(metadata_value, Mapping): + metadata = cast(Mapping[str, object], metadata_value) + else: + metadata = {} + user_id: object = ( + (user_api_key_dict.user_id if user_api_key_dict is not None else None) + or data.get("user") + or metadata.get("user_id") + or metadata.get("session_id") + or "anonymous" + ) + raw = "|".join( + ( + str(user_id), + str(call_type), + str(data.get("model", "")), + str(data.get("temperature", "")), + str(data.get("max_tokens", "")), + system_prompt, + last_user_prompt, + prompt, + ) + ) + return hashlib.sha256(raw.encode("utf-8")).hexdigest() + + def _content_text(self, content: object) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + content_items = cast(list[object], content) + return "\n".join( + text + for item in content_items + if isinstance(item, Mapping) + for text in [cast(Mapping[str, object], item).get("text")] + if isinstance(text, str) + ) + return "" + + +duplicate_burst_guard = DuplicateBurstGuard() diff --git a/tests/test_litellm/test_duplicate_burst_guard_example.py b/tests/test_litellm/test_duplicate_burst_guard_example.py new file mode 100644 index 00000000000..376a829fa3d --- /dev/null +++ b/tests/test_litellm/test_duplicate_burst_guard_example.py @@ -0,0 +1,80 @@ +from collections.abc import Mapping +from typing import cast + +import pytest +from fastapi import HTTPException + +from cookbook.litellm_proxy_server.duplicate_burst_guard import DuplicateBurstGuard +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral + +USER_API_KEY = UserAPIKeyAuth(user_id="user-1") +CACHE = DualCache() +CALL_TYPE: CallTypesLiteral = "completion" + + +def _chat_request(user_prompt: str) -> dict[str, object]: + return { + "model": "gpt-4o-mini", + "messages": [ + {"role": "system", "content": "Use the shared operating policy."}, + {"role": "user", "content": user_prompt}, + ], + "temperature": 0, + } + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_blocks_repeated_prompt() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request = _chat_request("Summarize invoice A") + + assert ( + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request, CALL_TYPE) + == request + ) + + with pytest.raises(HTTPException) as exc_info: + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request, CALL_TYPE) + + assert exc_info.value.status_code == 429 + detail = exc_info.value.detail + assert isinstance(detail, Mapping) + detail_data = cast(Mapping[str, object], detail) + assert detail_data.get("error") == "Duplicate request burst detected" + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_uses_last_user_prompt() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_b = _chat_request("Summarize invoice B") + + await guard.async_pre_call_hook( + USER_API_KEY, CACHE, _chat_request("Summarize invoice A"), CALL_TYPE + ) + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, CALL_TYPE + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_uses_completion_prompt() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a: dict[str, object] = { + "model": "gpt-4o-mini", + "prompt": "Summarize invoice A", + } + request_b: dict[str, object] = { + "model": "gpt-4o-mini", + "prompt": "Summarize invoice B", + } + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, CALL_TYPE) + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, CALL_TYPE + ) + + assert accepted == request_b From 30a6f1ae22fd831197a6121bff16283fe469e231 Mon Sep 17 00:00:00 2001 From: foxjl85-dev <272123326+foxjl85-dev@users.noreply.github.com> Date: Mon, 29 Jun 2026 15:41:29 -0500 Subject: [PATCH 2/3] docs(proxy): harden duplicate burst guard fingerprint --- .../duplicate_burst_guard.py | 94 +++++----- .../test_duplicate_burst_guard_example.py | 160 ++++++++++++++++++ 2 files changed, 209 insertions(+), 45 deletions(-) diff --git a/cookbook/litellm_proxy_server/duplicate_burst_guard.py b/cookbook/litellm_proxy_server/duplicate_burst_guard.py index c7d13b26ebb..33060687048 100644 --- a/cookbook/litellm_proxy_server/duplicate_burst_guard.py +++ b/cookbook/litellm_proxy_server/duplicate_burst_guard.py @@ -9,10 +9,15 @@ litellm_settings: This is an in-memory, single-process example. Use a shared store if your proxy runs multiple workers or needs duplicate detection across instances. + +The request fingerprint uses a whitelist of model-affecting fields and scopes +duplicates by request user/session before API-key owner. Adjust those fields if +your deployment needs different duplicate semantics. """ import asyncio import hashlib +import json import time from collections import defaultdict, deque from collections.abc import Mapping @@ -25,6 +30,34 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import CallTypesLiteral +FINGERPRINT_REQUEST_KEYS = ( + "model", + "messages", + "prompt", + "input", + "temperature", + "max_tokens", + "max_completion_tokens", + "top_p", + "stop", + "stream", + "tools", + "tool_choice", + "functions", + "function_call", + "response_format", + "n", + "presence_penalty", + "frequency_penalty", + "logit_bias", + "seed", + "modalities", + "audio", + "reasoning_effort", + "thinking", + "extra_body", +) + class DuplicateBurstGuard(CustomLogger): def __init__(self, max_calls: int = 2, window_seconds: float = 10.0) -> None: @@ -85,64 +118,35 @@ class DuplicateBurstGuard(CustomLogger): user_api_key_dict: Optional[UserAPIKeyAuth], call_type: CallTypesLiteral, ) -> str: - messages_value = data.get("messages") - system_prompt = "" - last_user_prompt = "" - prompt = "" - - if isinstance(messages_value, list): - messages = cast(list[object], messages_value) - for message in messages: - if not isinstance(message, Mapping): - continue - message_data = cast(Mapping[str, object], message) - if message_data.get("role") == "system" and not system_prompt: - system_prompt = self._content_text(message_data.get("content")) - if message_data.get("role") == "user": - last_user_prompt = self._content_text(message_data.get("content")) - else: - prompt = self._content_text(data.get("prompt")) - metadata_value = data.get("metadata") metadata: Mapping[str, object] if isinstance(metadata_value, Mapping): metadata = cast(Mapping[str, object], metadata_value) else: metadata = {} + key_user_id: object = None + if user_api_key_dict is not None: + key_user_id = user_api_key_dict.user_id user_id: object = ( - (user_api_key_dict.user_id if user_api_key_dict is not None else None) - or data.get("user") + data.get("user") or metadata.get("user_id") or metadata.get("session_id") + or key_user_id or "anonymous" ) - raw = "|".join( - ( - str(user_id), - str(call_type), - str(data.get("model", "")), - str(data.get("temperature", "")), - str(data.get("max_tokens", "")), - system_prompt, - last_user_prompt, - prompt, - ) + request = {key: data[key] for key in FINGERPRINT_REQUEST_KEYS if key in data} + raw = json.dumps( + { + "call_type": call_type, + "request": request, + "request_identity": user_id, + }, + default=str, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, ) return hashlib.sha256(raw.encode("utf-8")).hexdigest() - def _content_text(self, content: object) -> str: - if isinstance(content, str): - return content - if isinstance(content, list): - content_items = cast(list[object], content) - return "\n".join( - text - for item in content_items - if isinstance(item, Mapping) - for text in [cast(Mapping[str, object], item).get("text")] - if isinstance(text, str) - ) - return "" - duplicate_burst_guard = DuplicateBurstGuard() diff --git a/tests/test_litellm/test_duplicate_burst_guard_example.py b/tests/test_litellm/test_duplicate_burst_guard_example.py index 376a829fa3d..fa753f6005b 100644 --- a/tests/test_litellm/test_duplicate_burst_guard_example.py +++ b/tests/test_litellm/test_duplicate_burst_guard_example.py @@ -78,3 +78,163 @@ async def test_duplicate_burst_guard_uses_completion_prompt() -> None: ) assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_avoids_delimiter_collisions() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a = _chat_request("c") + request_a["messages"] = [ + {"role": "system", "content": "a|b"}, + {"role": "user", "content": "c"}, + ] + request_b = _chat_request("b|c") + request_b["messages"] = [ + {"role": "system", "content": "a"}, + {"role": "user", "content": "b|c"}, + ] + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, CALL_TYPE) + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, CALL_TYPE + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_multimodal_content() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a: dict[str, object] = { + "model": "gpt-4o-mini", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image_url", "image_url": {"url": "https://a.test/1.png"}}, + ], + } + ], + } + request_b: dict[str, object] = { + "model": "gpt-4o-mini", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image_url", "image_url": {"url": "https://a.test/2.png"}}, + ], + } + ], + } + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, CALL_TYPE) + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, CALL_TYPE + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_prefers_request_user_identity() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a = _chat_request("Summarize invoice A") + request_a["user"] = "end-user-a" + request_b = _chat_request("Summarize invoice A") + request_b["user"] = "end-user-b" + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, CALL_TYPE) + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, CALL_TYPE + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_full_chat_context() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a = _chat_request("What should I do next?") + request_a["messages"] = [ + {"role": "system", "content": "Use the shared operating policy."}, + {"role": "user", "content": "Invoice A is overdue."}, + {"role": "assistant", "content": "Ask for payment status."}, + {"role": "user", "content": "What should I do next?"}, + ] + request_b = _chat_request("What should I do next?") + request_b["messages"] = [ + {"role": "system", "content": "Use the shared operating policy."}, + {"role": "user", "content": "Invoice B has a credit memo."}, + {"role": "assistant", "content": "Check whether it was applied."}, + {"role": "user", "content": "What should I do next?"}, + ] + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, CALL_TYPE) + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, CALL_TYPE + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_output_options() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a = _chat_request("Summarize invoice A") + request_a["response_format"] = {"type": "json_object"} + request_b = _chat_request("Summarize invoice A") + request_b["tools"] = [ + { + "type": "function", + "function": { + "name": "summarize_invoice", + "parameters": {"type": "object"}, + }, + } + ] + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, CALL_TYPE) + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, CALL_TYPE + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_non_chat_input() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a: dict[str, object] = { + "model": "text-embedding-3-small", + "input": "invoice A", + } + request_b: dict[str, object] = { + "model": "text-embedding-3-small", + "input": "invoice B", + } + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, "aembedding") + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, "aembedding" + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_ignores_ephemeral_call_ids() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a = _chat_request("Summarize invoice A") + request_a["litellm_call_id"] = "call-a" + request_b = _chat_request("Summarize invoice A") + request_b["litellm_call_id"] = "call-b" + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, CALL_TYPE) + + with pytest.raises(HTTPException) as exc_info: + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_b, CALL_TYPE) + + assert exc_info.value.status_code == 429 From 67798ccd916ea9d064379d5c1bedbfe6476e4b7c Mon Sep 17 00:00:00 2001 From: foxjl85-dev <272123326+foxjl85-dev@users.noreply.github.com> Date: Mon, 29 Jun 2026 15:58:49 -0500 Subject: [PATCH 3/3] docs(proxy): cover duplicate guard review cases --- .../duplicate_burst_guard.py | 22 +++- .../test_duplicate_burst_guard_example.py | 123 ++++++++++++++++++ 2 files changed, 139 insertions(+), 6 deletions(-) diff --git a/cookbook/litellm_proxy_server/duplicate_burst_guard.py b/cookbook/litellm_proxy_server/duplicate_burst_guard.py index 33060687048..5170e7a3553 100644 --- a/cookbook/litellm_proxy_server/duplicate_burst_guard.py +++ b/cookbook/litellm_proxy_server/duplicate_burst_guard.py @@ -11,8 +11,8 @@ This is an in-memory, single-process example. Use a shared store if your proxy runs multiple workers or needs duplicate detection across instances. The request fingerprint uses a whitelist of model-affecting fields and scopes -duplicates by request user/session before API-key owner. Adjust those fields if -your deployment needs different duplicate semantics. +duplicates by authenticated API-key owner plus request user/session. Adjust +those fields if your deployment needs different duplicate semantics. """ import asyncio @@ -32,9 +32,14 @@ from litellm.types.utils import CallTypesLiteral FINGERPRINT_REQUEST_KEYS = ( "model", + "system", "messages", "prompt", "input", + "query", + "documents", + "instructions", + "previous_response_id", "temperature", "max_tokens", "max_completion_tokens", @@ -56,7 +61,9 @@ FINGERPRINT_REQUEST_KEYS = ( "reasoning_effort", "thinking", "extra_body", + "vector_store_id", ) +SKIP_CALL_TYPES: frozenset[str] = frozenset({"aimage_edit"}) class DuplicateBurstGuard(CustomLogger): @@ -76,6 +83,9 @@ class DuplicateBurstGuard(CustomLogger): data: dict[str, object], call_type: CallTypesLiteral, ) -> dict[str, object]: + if call_type in SKIP_CALL_TYPES: + return data + fingerprint = self._fingerprint( data=data, user_api_key_dict=user_api_key_dict, call_type=call_type ) @@ -127,19 +137,19 @@ class DuplicateBurstGuard(CustomLogger): key_user_id: object = None if user_api_key_dict is not None: key_user_id = user_api_key_dict.user_id - user_id: object = ( + request_user_id: object = ( data.get("user") or metadata.get("user_id") or metadata.get("session_id") - or key_user_id - or "anonymous" + or "anonymous_request" ) request = {key: data[key] for key in FINGERPRINT_REQUEST_KEYS if key in data} raw = json.dumps( { + "api_key_identity": key_user_id or "anonymous_key", "call_type": call_type, "request": request, - "request_identity": user_id, + "request_identity": request_user_id, }, default=str, ensure_ascii=False, diff --git a/tests/test_litellm/test_duplicate_burst_guard_example.py b/tests/test_litellm/test_duplicate_burst_guard_example.py index fa753f6005b..75d5672ca77 100644 --- a/tests/test_litellm/test_duplicate_burst_guard_example.py +++ b/tests/test_litellm/test_duplicate_burst_guard_example.py @@ -10,6 +10,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import CallTypesLiteral USER_API_KEY = UserAPIKeyAuth(user_id="user-1") +OTHER_USER_API_KEY = UserAPIKeyAuth(user_id="user-2") CACHE = DualCache() CALL_TYPE: CallTypesLiteral = "completion" @@ -238,3 +239,125 @@ async def test_duplicate_burst_guard_ignores_ephemeral_call_ids() -> None: await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_b, CALL_TYPE) assert exc_info.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_authenticated_key_identity() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request = _chat_request("Summarize invoice A") + request["metadata"] = {"session_id": "shared-session"} + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request, CALL_TYPE) + accepted = await guard.async_pre_call_hook( + OTHER_USER_API_KEY, CACHE, request, CALL_TYPE + ) + + assert accepted == request + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_rerank_fields() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a: dict[str, object] = { + "model": "rerank-english-v3.0", + "query": "invoice A", + "documents": ["invoice A is overdue"], + } + request_b: dict[str, object] = { + "model": "rerank-english-v3.0", + "query": "invoice B", + "documents": ["invoice B has a credit memo"], + } + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, "rerank") + accepted = await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_b, "rerank") + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_responses_context() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a: dict[str, object] = { + "model": "gpt-4o-mini", + "input": "continue", + "previous_response_id": "resp-a", + "instructions": "Use policy A", + } + request_b: dict[str, object] = { + "model": "gpt-4o-mini", + "input": "continue", + "previous_response_id": "resp-b", + "instructions": "Use policy B", + } + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, "aresponses") + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, "aresponses" + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_top_level_system_prompt() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a: dict[str, object] = { + "model": "claude-sonnet-4-5", + "system": "Use policy A", + "messages": [{"role": "user", "content": "Summarize invoice A"}], + } + request_b: dict[str, object] = { + "model": "claude-sonnet-4-5", + "system": "Use policy B", + "messages": [{"role": "user", "content": "Summarize invoice A"}], + } + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, CALL_TYPE) + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, CALL_TYPE + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_includes_search_fields() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request_a: dict[str, object] = { + "model": "search-model", + "query": "invoice A", + "vector_store_id": "store-a", + } + request_b: dict[str, object] = { + "model": "search-model", + "query": "invoice B", + "vector_store_id": "store-b", + } + + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request_a, "asearch") + accepted = await guard.async_pre_call_hook( + USER_API_KEY, CACHE, request_b, "asearch" + ) + + assert accepted == request_b + + +@pytest.mark.asyncio +async def test_duplicate_burst_guard_skips_image_edits() -> None: + guard = DuplicateBurstGuard(max_calls=1, window_seconds=60) + request: dict[str, object] = { + "model": "gpt-image-1", + "prompt": "Add the logo", + "image": object(), + "mask": object(), + } + + assert ( + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request, "aimage_edit") + == request + ) + assert ( + await guard.async_pre_call_hook(USER_API_KEY, CACHE, request, "aimage_edit") + == request + )