From febb718f216f6b20f334ca81c4a11b28bd863dc0 Mon Sep 17 00:00:00 2001 From: Surender Gulia Date: Fri, 30 Oct 2026 11:42:24 +0000 Subject: [PATCH 1/6] fix(proxy/batches): attribute end_user_id to completed batch spend logs (closes #49202) --- .../proxy/common_utils/managed_object_attribution.py | 8 ++++++-- .../litellm_enterprise/proxy/hooks/managed_files.py | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index a4a40a7a3d7..addfc99e47a 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -159,8 +159,12 @@ class ManagedObjectCreatorAttribution: async def batch_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]: rows: Final = await self._rows(job, object_id) - return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)} - + end_user_id: Final = getattr(job, "end_user_id", None) + return { + "user_api_key_user_id": job.created_by, + **_shared_metadata(job, rows), + **({"user_api_key_end_user_id": end_user_id} if end_user_id else {}), + } async def response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator: try: return await self._read_response_creator(job, object_id) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5b54e7d72a0..21ecdbe2d72 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -736,7 +736,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): from prisma import Json api_key = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None - end_user_id: Final = user_api_key_dict.end_user_id if file_purpose == "response" else None + end_user_id: Final = user_api_key_dict.end_user_id if file_purpose in ["response", "batch"] else None attribution_columns = ( { **({"api_key": api_key} if api_key is not None else {}), From 18753c39625e491d93027ecb7e6726481cbfdb93 Mon Sep 17 00:00:00 2001 From: Surender Gulia Date: Mon, 2 Nov 2026 07:16:44 +0000 Subject: [PATCH 2/6] test(proxy/batches): add regression test for end-user budget enforcement (closes #49202) --- .../spend/test_batch_completion_accounting.py | 104 ++++++++++++++++++ 1 file changed, 104 insertions(+) diff --git a/tests/integration/spend/test_batch_completion_accounting.py b/tests/integration/spend/test_batch_completion_accounting.py index 6973864117c..5d5e40b96ed 100644 --- a/tests/integration/spend/test_batch_completion_accounting.py +++ b/tests/integration/spend/test_batch_completion_accounting.py @@ -430,3 +430,107 @@ def test_batch_reconciliation_delivers_each_configured_logger_once(gateway: Gate BATCH_PROMPT_TOKENS, BATCH_COMPLETION_TOKENS, ) + + +def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> None: + """ + Regression test for #49202: + Validates end-to-end attribution and budget enforcement for batch operations. + + Verifies: + 1. Header attribution: 'x-litellm-end-user-id' propagates through batch creation + and CheckBatchCost settlement to update LiteLLM_EndUserTable.spend. + 2. Budget enforcement: Once batch spend exceeds customer max_budget, subsequent + requests for customer A are denied with budget exceeded. + 3. Identity isolation: Customer B on the same API key remains unblocked and unaffected. + 4. Observation immutability: Subsequent reads of the batch do not mutate stored creator attribution. + """ + with gateway.scenario() as scenario: + scenario_id: Final = f"batch-end-user-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario(scenario_id, batch_routes("gpt-4o-mini")) + scenario.cleanups.callback(delete_scenario, handle) + + model: Final = scenario.model( + api_base=handle.api_base(), + input_cost_per_token=INPUT_COST_PER_TOKEN, + output_cost_per_token=OUTPUT_COST_PER_TOKEN, + ) + key: Final = scenario.key(models=[model]) + + # 1. Setup Customer A with a low budget that the batch will exhaust + user_a: Final = f"user-a-{uuid.uuid4().hex[:8]}" + gateway.post("/customer/new", {"user_id": user_a, "max_budget": 0.001}) + + # 2. Setup Customer B on the same key with sufficient budget for isolation check + user_b: Final = f"user-b-{uuid.uuid4().hex[:8]}" + gateway.post("/customer/new", {"user_id": user_b, "max_budget": 10.0}) + + # 3. Upload batch input file + file_response: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "model": model}, + {"file": ("in.jsonl", batch_input_file(model), "application/jsonl")}, + key=key, + ) + assert file_response.status_code == 200, file_response.text + file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"]) + + # 4. Create batch with Customer A's end_user_id header + batch_response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": file_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + key=key, + headers={"x-litellm-end-user-id": user_a}, + ) + assert batch_response.status_code == 200, batch_response.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + + # 5. Observation check: Fetch batch to confirm terminal completion + retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key) + assert retrieval.status_code == 200, retrieval.text + assert retrieval.json()["status"] == "completed", retrieval.text + + # 6. Verify CheckBatchCost settlement: LiteLLM_EndUserTable reflects Customer A's spend + end_user_rows: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', + (user_a,), + ), + lambda values: len(values) == 1 and float(str(values[0]["spend"])) > 0.001, + seconds=70, + ) + assert len(end_user_rows) == 1, f"Expected spend row for {user_a}" + + # 7. Positive enforcement: Customer A's subsequent call is denied due to exhausted budget + denied_response: Final = gateway.request( + "POST", + "/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "ping"}], + }, + key=key, + headers={"x-litellm-end-user-id": user_a}, + ) + assert denied_response.status_code in [400, 422], denied_response.text + assert "budget" in denied_response.text.lower(), denied_response.text + + # 8. Identity isolation: Customer B remains unaffected on the same key + isolated_b_response: Final = gateway.request( + "POST", + "/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "ping"}], + }, + key=key, + headers={"x-litellm-end-user-id": user_b}, + ) + # Customer B is not blocked by Customer A's exhausted budget + assert isolated_b_response.status_code != denied_response.status_code or "budget" not in isolated_b_response.text.lower() \ No newline at end of file From 8b18c26eba8a436e84c780aae1cb710e6385daf5 Mon Sep 17 00:00:00 2001 From: Surender Gulia Date: Mon, 2 Nov 2026 09:14:51 +0000 Subject: [PATCH 3/6] test(proxy/batches): cover B-batch isolation and replay immutability (closes #49202) --- .../spend/test_batch_completion_accounting.py | 120 +++++++++++++----- 1 file changed, 87 insertions(+), 33 deletions(-) diff --git a/tests/integration/spend/test_batch_completion_accounting.py b/tests/integration/spend/test_batch_completion_accounting.py index 5d5e40b96ed..b73fed6827c 100644 --- a/tests/integration/spend/test_batch_completion_accounting.py +++ b/tests/integration/spend/test_batch_completion_accounting.py @@ -435,15 +435,17 @@ def test_batch_reconciliation_delivers_each_configured_logger_once(gateway: Gate def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> None: """ Regression test for #49202: - Validates end-to-end attribution and budget enforcement for batch operations. + Validates end-to-end attribution, budget enforcement, identity isolation, + and observation immutability for asynchronous batch operations. Verifies: - 1. Header attribution: 'x-litellm-end-user-id' propagates through batch creation - and CheckBatchCost settlement to update LiteLLM_EndUserTable.spend. - 2. Budget enforcement: Once batch spend exceeds customer max_budget, subsequent - requests for customer A are denied with budget exceeded. - 3. Identity isolation: Customer B on the same API key remains unblocked and unaffected. - 4. Observation immutability: Subsequent reads of the batch do not mutate stored creator attribution. + 1. Positive Attribution & Enforcement: Customer A's batch spend settles to + LiteLLM_EndUserTable, and subsequent calls for A are denied with budget_exceeded. + 2. Identity Isolation: Customer B on the same API key runs a separate batch; + its charge settles to B's record without cross-contaminating A, and B remains unblocked. + 3. Creation vs Observation Immutability: Retrieving Customer A's completed batch + using a different key does not overwrite the stored creator attribution. + 4. Replay Idempotency: Subsequent retrievals / poll cycles do not double-bill Customer A. """ with gateway.scenario() as scenario: scenario_id: Final = f"batch-end-user-{uuid.uuid4().hex[:12]}" @@ -456,31 +458,31 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No output_cost_per_token=OUTPUT_COST_PER_TOKEN, ) key: Final = scenario.key(models=[model]) + viewer_key: Final = scenario.key(models=[model]) - # 1. Setup Customer A with a low budget that the batch will exhaust + # 1. Setup Customer A (low budget: 0.001) and Customer B (ample budget: 10.0) user_a: Final = f"user-a-{uuid.uuid4().hex[:8]}" gateway.post("/customer/new", {"user_id": user_a, "max_budget": 0.001}) - # 2. Setup Customer B on the same key with sufficient budget for isolation check user_b: Final = f"user-b-{uuid.uuid4().hex[:8]}" gateway.post("/customer/new", {"user_id": user_b, "max_budget": 10.0}) - # 3. Upload batch input file - file_response: Final = gateway.request_multipart( + # 2. Upload batch input file for Customer A + file_a_resp: Final = gateway.request_multipart( "/v1/files", {"purpose": "batch", "model": model}, {"file": ("in.jsonl", batch_input_file(model), "application/jsonl")}, key=key, ) - assert file_response.status_code == 200, file_response.text - file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"]) + assert file_a_resp.status_code == 200, file_a_resp.text + file_a_id: Final = string_value(JSON_OBJECT.validate_json(file_a_resp.content)["id"]) - # 4. Create batch with Customer A's end_user_id header - batch_response: Final = gateway.request( + # 3. Create batch for Customer A + batch_a_resp: Final = gateway.request( "POST", "/v1/batches", { - "input_file_id": file_id, + "input_file_id": file_a_id, "endpoint": "/v1/chat/completions", "completion_window": "24h", "model": model, @@ -488,16 +490,16 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No key=key, headers={"x-litellm-end-user-id": user_a}, ) - assert batch_response.status_code == 200, batch_response.text - batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + assert batch_a_resp.status_code == 200, batch_a_resp.text + batch_a_id: Final = string_value(JSON_OBJECT.validate_json(batch_a_resp.content)["id"]) - # 5. Observation check: Fetch batch to confirm terminal completion - retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key) - assert retrieval.status_code == 200, retrieval.text - assert retrieval.json()["status"] == "completed", retrieval.text + # 4. Wait for Customer A batch completion + retrieval_a: Final = gateway.request("GET", f"/v1/batches/{batch_a_id}", key=key) + assert retrieval_a.status_code == 200, retrieval_a.text + assert retrieval_a.json()["status"] == "completed", retrieval_a.text - # 6. Verify CheckBatchCost settlement: LiteLLM_EndUserTable reflects Customer A's spend - end_user_rows: Final = eventually( + # 5. Verify CheckBatchCost settlement for Customer A in LiteLLM_EndUserTable + user_a_spend_rows: Final = eventually( lambda: read_rows( 'SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,), @@ -505,10 +507,11 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No lambda values: len(values) == 1 and float(str(values[0]["spend"])) > 0.001, seconds=70, ) - assert len(end_user_rows) == 1, f"Expected spend row for {user_a}" + spend_a_initial: Final = float(str(user_a_spend_rows[0]["spend"])) + assert spend_a_initial > 0.001, f"Expected spend for {user_a}" - # 7. Positive enforcement: Customer A's subsequent call is denied due to exhausted budget - denied_response: Final = gateway.request( + # 6. Customer A subsequent request must be denied due to exceeded budget + denied_a: Final = gateway.request( "POST", "/chat/completions", { @@ -518,11 +521,55 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No key=key, headers={"x-litellm-end-user-id": user_a}, ) - assert denied_response.status_code in [400, 422], denied_response.text - assert "budget" in denied_response.text.lower(), denied_response.text + assert denied_a.status_code in [400, 422], denied_a.text + assert "budget" in denied_a.text.lower(), denied_a.text - # 8. Identity isolation: Customer B remains unaffected on the same key - isolated_b_response: Final = gateway.request( + # 7. Customer B creates and settles their own batch using the SAME API key + file_b_resp: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "model": model}, + {"file": ("in.jsonl", batch_input_file(model), "application/jsonl")}, + key=key, + ) + assert file_b_resp.status_code == 200, file_b_resp.text + file_b_id: Final = string_value(JSON_OBJECT.validate_json(file_b_resp.content)["id"]) + + batch_b_resp: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": file_b_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + key=key, + headers={"x-litellm-end-user-id": user_b}, + ) + assert batch_b_resp.status_code == 200, batch_b_resp.text + batch_b_id: Final = string_value(JSON_OBJECT.validate_json(batch_b_resp.content)["id"]) + + retrieval_b: Final = gateway.request("GET", f"/v1/batches/{batch_b_id}", key=key) + assert retrieval_b.status_code == 200, retrieval_b.text + assert retrieval_b.json()["status"] == "completed", retrieval_b.text + + # 8. Verify Customer B's batch settled to B's record + user_b_spend_rows: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', + (user_b,), + ), + lambda values: len(values) == 1 and float(str(values[0]["spend"])) > 0, + seconds=70, + ) + assert len(user_b_spend_rows) == 1 + + # Assert no cross-contamination: Customer A's spend did not absorb B's batch + spend_a_current: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) + assert spend_a_current == spend_a_initial, "Customer A spend must not change when Customer B runs a batch" + + # Customer B is NOT budget-denied (has ample budget) + b_request: Final = gateway.request( "POST", "/chat/completions", { @@ -532,5 +579,12 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No key=key, headers={"x-litellm-end-user-id": user_b}, ) - # Customer B is not blocked by Customer A's exhausted budget - assert isolated_b_response.status_code != denied_response.status_code or "budget" not in isolated_b_response.text.lower() \ No newline at end of file + assert b_request.status_code != denied_a.status_code or "budget" not in b_request.text.lower() + + # 9. Observation Immutability: Retrieve Customer A's batch as another authorized key + retrieval_by_viewer: Final = gateway.request("GET", f"/v1/batches/{batch_a_id}", key=viewer_key) + assert retrieval_by_viewer.status_code == 200, retrieval_by_viewer.text + + # 10. Replay Control: Verify single billing idempotency (no double-counting) + spend_a_after_replay: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) + assert spend_a_after_replay == spend_a_initial, "Replay/retrieval must not double-bill customer A" \ No newline at end of file From 813a197bb6537f4077ef39ccecebb0103c7460de Mon Sep 17 00:00:00 2001 From: Surender Gulia Date: Mon, 2 Nov 2026 10:12:29 +0000 Subject: [PATCH 4/6] test(proxy/batches): assert ManagedObjectTable creator immutability and single-billing replay (closes #49202) --- .../spend/test_batch_completion_accounting.py | 41 +++++++++++++++++-- 1 file changed, 37 insertions(+), 4 deletions(-) diff --git a/tests/integration/spend/test_batch_completion_accounting.py b/tests/integration/spend/test_batch_completion_accounting.py index b73fed6827c..4a5cb36cc7f 100644 --- a/tests/integration/spend/test_batch_completion_accounting.py +++ b/tests/integration/spend/test_batch_completion_accounting.py @@ -581,10 +581,43 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No ) assert b_request.status_code != denied_a.status_code or "budget" not in b_request.text.lower() - # 9. Observation Immutability: Retrieve Customer A's batch as another authorized key + # 9. Creation vs Observation Immutability: Assert stored creator attribution before and after viewer read + creator_before: Final = read_rows( + 'SELECT end_user_id FROM "LiteLLM_ManagedObjectTable" WHERE object_id=%s', + (batch_a_id,), + )[0]["end_user_id"] + assert creator_before == user_a, f"Expected creator {user_a}, got {creator_before}" + + # Retrieve Customer A's batch using an alternate authorized key retrieval_by_viewer: Final = gateway.request("GET", f"/v1/batches/{batch_a_id}", key=viewer_key) assert retrieval_by_viewer.status_code == 200, retrieval_by_viewer.text - # 10. Replay Control: Verify single billing idempotency (no double-counting) - spend_a_after_replay: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) - assert spend_a_after_replay == spend_a_initial, "Replay/retrieval must not double-bill customer A" \ No newline at end of file + # Assert creator attribution remains intact and unmutated by observation + creator_after: Final = read_rows( + 'SELECT end_user_id FROM "LiteLLM_ManagedObjectTable" WHERE object_id=%s', + (batch_a_id,), + )[0]["end_user_id"] + assert creator_after == user_a, "Observation must not overwrite stored creator end_user_id" + + # 10. Replay / Single Billing Control: Verify spend row is exactly once and A remains over budget + a_spend_logs: Final = read_rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND call_type=\'aretrieve_batch\'', + (sha256(key.encode()).hexdigest(),), + ) + assert len(a_spend_logs) >= 1 + current_a_spend: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) + assert current_a_spend == spend_a_initial, "Customer A must be billed exactly once" + + # Verify Customer A remains strictly budget-denied + recheck_denied: Final = gateway.request( + "POST", + "/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "ping"}], + }, + key=key, + headers={"x-litellm-end-user-id": user_a}, + ) + assert recheck_denied.status_code in [400, 422], recheck_denied.text + assert "budget" in recheck_denied.text.lower(), recheck_denied.text \ No newline at end of file From 3c3e422f11cfe6bad712286596c17367f8161ed5 Mon Sep 17 00:00:00 2001 From: Surender Gulia Date: Mon, 2 Nov 2026 11:25:26 +0000 Subject: [PATCH 5/6] test(proxy/batches): enforce unified_object_id schema and strictly assert settlement idempotency --- .../spend/test_batch_completion_accounting.py | 42 +++++++++++++++---- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/tests/integration/spend/test_batch_completion_accounting.py b/tests/integration/spend/test_batch_completion_accounting.py index 4a5cb36cc7f..2db4ce96f08 100644 --- a/tests/integration/spend/test_batch_completion_accounting.py +++ b/tests/integration/spend/test_batch_completion_accounting.py @@ -583,7 +583,7 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No # 9. Creation vs Observation Immutability: Assert stored creator attribution before and after viewer read creator_before: Final = read_rows( - 'SELECT end_user_id FROM "LiteLLM_ManagedObjectTable" WHERE object_id=%s', + 'SELECT end_user_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id=%s', (batch_a_id,), )[0]["end_user_id"] assert creator_before == user_a, f"Expected creator {user_a}, got {creator_before}" @@ -592,22 +592,50 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No retrieval_by_viewer: Final = gateway.request("GET", f"/v1/batches/{batch_a_id}", key=viewer_key) assert retrieval_by_viewer.status_code == 200, retrieval_by_viewer.text - # Assert creator attribution remains intact and unmutated by observation + # Invariant 3: Unified Object Table creator immutability creator_after: Final = read_rows( - 'SELECT end_user_id FROM "LiteLLM_ManagedObjectTable" WHERE object_id=%s', + 'SELECT end_user_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id=%s', (batch_a_id,), )[0]["end_user_id"] assert creator_after == user_a, "Observation must not overwrite stored creator end_user_id" - # 10. Replay / Single Billing Control: Verify spend row is exactly once and A remains over budget + # 10. Replay / Single Billing Control: Scoped to Batch A settlement identity (Exactly-Once) a_spend_logs: Final = read_rows( - 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND call_type=\'aretrieve_batch\'', - (sha256(key.encode()).hexdigest(),), + 'SELECT spend, end_user FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND call_type=\'aretrieve_batch\' AND end_user=%s', + (sha256(key.encode()).hexdigest(), user_a), ) - assert len(a_spend_logs) >= 1 + assert len(a_spend_logs) == 1, f"Expected exactly 1 spend log for Batch A, got {len(a_spend_logs)}" + + # Verify batch_processed is marked True on ManagedObjectTable + batch_a_processed: Final = read_rows( + 'SELECT batch_processed FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id=%s', + (batch_a_id,), + )[0]["batch_processed"] + assert bool(batch_a_processed) is True, "batch_processed must be True after completion" + current_a_spend: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) assert current_a_spend == spend_a_initial, "Customer A must be billed exactly once" + # --- EXPLICIT SECOND RE-POLL PASS (IDEMPOTENCY VERIFICATION) --- + # Re-trigger batch retrieve to ensure no duplicate CheckBatchCost debit occurs + re_retrieve = gateway.request( + "GET", + f"/batches/{batch_a_id}", + key=key, + ) + assert re_retrieve.status_code == 200, re_retrieve.text + + # Verify spend logs count remains strictly 1 (no double spend row appended) + recheck_spend_logs: Final = read_rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND call_type=\'aretrieve_batch\' AND end_user=%s', + (sha256(key.encode()).hexdigest(), user_a), + ) + assert len(recheck_spend_logs) == 1, "Duplicate spend log detected after re-poll" + + # Verify EndUserTable spend has not doubled + post_repoll_spend: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) + assert post_repoll_spend == current_a_spend, f"Double charge detected! {post_repoll_spend} != {current_a_spend}" + # Verify Customer A remains strictly budget-denied recheck_denied: Final = gateway.request( "POST", From c7bc0c60458a50d77dcab56cba564c1e12c0010b Mon Sep 17 00:00:00 2001 From: Surender Gulia Date: Mon, 2 Nov 2026 16:22:23 +0000 Subject: [PATCH 6/6] test(proxy/batches): label GET assertions as post-settlement observation idempotency --- .../spend/test_batch_completion_accounting.py | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/integration/spend/test_batch_completion_accounting.py b/tests/integration/spend/test_batch_completion_accounting.py index 2db4ce96f08..d8395daae7c 100644 --- a/tests/integration/spend/test_batch_completion_accounting.py +++ b/tests/integration/spend/test_batch_completion_accounting.py @@ -10,11 +10,11 @@ from typing import Final import pytest import yaml -from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value -from integration._support.database import read_rows, scratch_database -from integration._support.process import owned_proxy -from integration._support.wire import Reply, Request, wire_server -from integration._support.upstream import delete_scenario, register_scenario +from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from tests.integration._support.database import read_rows, scratch_database +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, wire_server +from tests.integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse from pydantic import JsonValue @@ -616,25 +616,25 @@ def test_batch_end_user_budget_enforcement_and_isolation(gateway: Gateway) -> No current_a_spend: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) assert current_a_spend == spend_a_initial, "Customer A must be billed exactly once" - # --- EXPLICIT SECOND RE-POLL PASS (IDEMPOTENCY VERIFICATION) --- - # Re-trigger batch retrieve to ensure no duplicate CheckBatchCost debit occurs - re_retrieve = gateway.request( + # --- POST-SETTLEMENT OBSERVATION IDEMPOTENCY --- + # Verify that subsequent read/GET requests on a terminal batch do not trigger billing side effects + re_retrieve: Final = gateway.request( "GET", f"/batches/{batch_a_id}", key=key, ) assert re_retrieve.status_code == 200, re_retrieve.text - # Verify spend logs count remains strictly 1 (no double spend row appended) - recheck_spend_logs: Final = read_rows( + # Verify spend logs count remains strictly 1 (GET observation does not append spend rows) + post_observation_spend_logs: Final = read_rows( 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND call_type=\'aretrieve_batch\' AND end_user=%s', (sha256(key.encode()).hexdigest(), user_a), ) - assert len(recheck_spend_logs) == 1, "Duplicate spend log detected after re-poll" + assert len(post_observation_spend_logs) == 1, "Observation side-effect: duplicate spend row appended on terminal batch GET" - # Verify EndUserTable spend has not doubled - post_repoll_spend: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) - assert post_repoll_spend == current_a_spend, f"Double charge detected! {post_repoll_spend} != {current_a_spend}" + # Verify EndUserTable spend has not incremented + post_observation_spend: Final = float(str(read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_a,))[0]["spend"])) + assert post_observation_spend == current_a_spend, f"Observation side-effect: spend increased from {current_a_spend} to {post_observation_spend}" # Verify Customer A remains strictly budget-denied recheck_denied: Final = gateway.request(