From d5531234ce9bd37de13bf00ef9d98debca8dc62a Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 15 Oct 2026 04:08:56 -0700
Subject: [PATCH] test(ci): pin xai reported-cost split and coordinator-held
 Prisma lock timeout (#46645)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 tests/e2e/migrations/database.py                 | 9 +++++++++
 tests/e2e/migrations/test_startup.py             | 9 ++++++---
 tests/integration/cost_calculation/cases/misc.py | 8 +++++++-
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/tests/e2e/migrations/database.py b/tests/e2e/migrations/database.py
index 9bf03a216c4..34d94025a6e 100644
--- a/tests/e2e/migrations/database.py
+++ b/tests/e2e/migrations/database.py
@@ -89,6 +89,15 @@ class Database:
             (key >> 32, key & 0xFFFFFFFF),
         )
 
+    @step("List the database sessions holding an advisory lock")
+    def holders(self, key: int = GATE_KEY) -> tuple[tuple[Scalar, ...], ...]:
+        return self.query(
+            "SELECT pid FROM pg_locks WHERE locktype = 'advisory' AND granted "
+            "AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) "
+            "AND classid = %s AND objid = %s ORDER BY pid",
+            (key >> 32, key & 0xFFFFFFFF),
+        )
+
     @step("List replica identities waiting on any of the advisory locks")
     def waiting_identities(self, keys: tuple[int, ...]) -> frozenset[str]:
         if not keys:
diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py
index 4a43ed4c0e3..5ec8d7690f2 100644
--- a/tests/e2e/migrations/test_startup.py
+++ b/tests/e2e/migrations/test_startup.py
@@ -14,7 +14,7 @@ from .checks import (
     start_replicas,
 )
 from .containers import Containers, failed, ready, until, waiting
-from .database import GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user
+from .database import COORDINATOR_LOCK, GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user
 from .startup_models import Migration
 from e2e_metadata import Domain, Subject, meta
 
@@ -97,8 +97,11 @@ class TestMigrationStartup:
             with containers.start(
                 database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "12"}
             ) as replica:
-                until("Prisma lock contention", lambda: bool(database.blocked(PRISMA_LOCK)))
-                failed((replica,), "Timed out waiting for")
+                until(
+                    "a replica holding the coordinator lock while Prisma's lock is held",
+                    lambda: bool(database.holders(COORDINATOR_LOCK)),
+                )
+                failed((replica,), "Timed out waiting for Prisma's lock")
                 assert database.history() == history
                 assert not database.exists("migration_effect")
         with containers.start(database, (COMPLETE,)) as restarted:
diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py
index 389e552e44c..d17ab7a3d31 100644
--- a/tests/integration/cost_calculation/cases/misc.py
+++ b/tests/integration/cost_calculation/cases/misc.py
@@ -599,7 +599,13 @@ XAI_GROK_5_PROVIDER_REPORTED_COST: Final = CostTrackingTestCase(
             "usage": {"prompt_tokens": 1840, "completion_tokens": 412, "total_tokens": 2252, "cost": 0.0421},
         },
     ),
-    expected=ExactExpected(spend=0.0421, input_cost=0.0, output_cost=0.0421, prompt_tokens=1840, completion_tokens=412),
+    expected=ExactExpected(
+        spend=0.0421,
+        input_cost=0.029078078,
+        output_cost=0.013021922,
+        prompt_tokens=1840,
+        completion_tokens=412,
+    ),
 )