From f07a7cbcde623f9b33323641842718ececb5afd8 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Mon, 13 Jul 2026 17:26:10 -0400
Subject: [PATCH] fix(e2e): read spend logs via paginated /spend/logs/v2 to
 stop OOM'ing the runner

The batches suite's test_rate_limited_batch_create_leaves_no_unattributed_spend_row
snapshots every unattributed spend row twice through spend_logs(SpendLogsParams()),
which used the deprecated, unpaginated /spend/logs. With no filter that endpoint
runs find_all over the whole LiteLLM_SpendLogs table, heavy message/response
columns included, so against a shared staging database the response OOM-killed the
e2e pod mid-test and the run restarted from preflight every time, always dying at
that same test

Repoint Gateway.spend_logs at /spend/logs/v2, which is paginated, drops the heavy
columns, and caps the row count. It walks every page over a bounded default window
and returns the same list[SpendLogRow], so the ~15 poll_logs_for_key / spend_logs
callers are untouched. v2 requires an explicit date window and matches the token
as stored (a SHA-256 hash) rather than the raw sk- key, so spend_logs fills a
bounded window when the caller gives none and hashes an sk- key the way the proxy
hashes it; the window and pagination fields are added to SpendLogsParams
---
 tests/e2e/e2e_gateway.py | 48 ++++++++++++++++++++++++++++++++++------
 tests/e2e/models.py      |  4 ++++
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/tests/e2e/e2e_gateway.py b/tests/e2e/e2e_gateway.py
index 05f83ecc085..93f45d18c60 100644
--- a/tests/e2e/e2e_gateway.py
+++ b/tests/e2e/e2e_gateway.py
@@ -8,10 +8,12 @@ Gateway's key/customer methods for cleanup. Read-backs are eventually consistent
 
 from __future__ import annotations
 
+import hashlib
 import time
 import warnings
 from collections.abc import Callable
 from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
 
 from e2e_http import (
     NoBody,
@@ -49,7 +51,7 @@ from models import (
     OcrBody,
     OcrResponse,
     SpendLogRow,
-    SpendLogs,
+    SpendLogsPage,
     SpendLogsParams,
 )
 from e2e_config import (
@@ -65,6 +67,15 @@ from transport import HttpTransport, SplitTransport, Transport
 RowsPredicate = Callable[[list[SpendLogRow]], bool]
 
 
+def _hashed_if_raw_key(api_key: str | None) -> str | None:
+    """The v2 spend-logs filter matches the token as stored on the row (a
+    SHA-256 hash), not the raw ``sk-`` key. Mirror the proxy's ``hash_token`` so
+    a raw key filters correctly; an already-hashed value passes through."""
+    if api_key is None or not api_key.startswith("sk-"):
+        return api_key
+    return hashlib.sha256(api_key.encode()).hexdigest()
+
+
 @dataclass(frozen=True, slots=True)
 class Gateway:
     transport: Transport
@@ -243,17 +254,40 @@ class Gateway:
     # ---- spend read-back ------------------------------------------------
 
     def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]:
+        """Every spend row matching ``params``, read by paging through the
+        paginated /spend/logs/v2. This replaces the deprecated, unpaginated
+        /spend/logs, which full-table-scans the entire spend history with the
+        heavy message/response columns included and OOM-kills the caller on a
+        large database. v2 needs an explicit window and matches the hashed token
+        the proxy stores, so a bounded default window is filled when the caller
+        gives none and an ``sk-`` key is hashed the way the proxy hashes it."""
+        now = datetime.now(timezone.utc)
+        fmt = "%Y-%m-%d %H:%M:%S"
+        query = params.model_copy(
+            update={
+                "api_key": _hashed_if_raw_key(params.api_key),
+                "start_date": params.start_date or (now - timedelta(days=1)).strftime(fmt),
+                "end_date": params.end_date or (now + timedelta(days=1)).strftime(fmt),
+            }
+        )
+        first = self._spend_logs_page(query, page=1)
+        if first is None:
+            return []
+        rest = (self._spend_logs_page(query, page=page) for page in range(2, first.total_pages + 1))
+        return [*first.data, *(row for page in rest if page is not None for row in page.data)]
+
+    def _spend_logs_page(self, query: SpendLogsParams, *, page: int) -> SpendLogsPage | None:
         result = self.transport.get(
-            "/spend/logs",
+            "/spend/logs/v2",
             headers=self.transport.master,
-            params=params,
-            response_type=SpendLogs,
+            params=query.model_copy(update={"page": page}),
+            response_type=SpendLogsPage,
         )
         match result:
-            case Success(data=logs):
-                return logs.root
+            case Success(data=page_data):
+                return page_data
             case _:
-                return []
+                return None
 
     def poll_logs_for_key(
         self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None
diff --git a/tests/e2e/models.py b/tests/e2e/models.py
index a977f2de542..f7717a86ea5 100644
--- a/tests/e2e/models.py
+++ b/tests/e2e/models.py
@@ -274,6 +274,10 @@ class SpendLogs(RootModel[list[SpendLogRow]]):
 class SpendLogsParams(BaseModel):
     request_id: str | None = None
     api_key: str | None = None
+    start_date: str | None = None
+    end_date: str | None = None
+    page: int = 1
+    page_size: int = 100
 
 
 class SpendLogsPageParams(BaseModel):