From e66981c7e86343639c705ea0c8d77460c5a75311 Mon Sep 17 00:00:00 2001
From: Shubham Naik <shub@letta.com>
Date: Mon, 16 Feb 2026 15:07:06 -0800
Subject: [PATCH] feat: update undertaker to use rate limiter (#9498)

---
 letta/agents/letta_agent_v2.py                |  9 +++++++--
 letta/services/credit_verification_service.py | 16 ++++++++++------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/letta/agents/letta_agent_v2.py b/letta/agents/letta_agent_v2.py
index 02ed2e46..e5bdec0f 100644
--- a/letta/agents/letta_agent_v2.py
+++ b/letta/agents/letta_agent_v2.py
@@ -705,10 +705,15 @@ class LettaAgentV2(BaseAgentV2):
     async def _check_credits(self) -> bool:
         """Check if the organization still has credits. Returns True if OK or not configured."""
         try:
-            await self.credit_verification_service.verify_credits(self.actor.organization_id)
+            await self.credit_verification_service.verify_credits(
+                self.actor.organization_id, self.agent_state.id
+            )
             return True
         except InsufficientCreditsError:
-            self.logger.warning(f"Insufficient credits for organization {self.actor.organization_id}, stopping agent loop")
+            self.logger.warning(
+                f"Insufficient credits for organization {self.actor.organization_id}, "
+                f"agent {self.agent_state.id}, stopping agent loop"
+            )
             return False
 
     @trace_method
diff --git a/letta/services/credit_verification_service.py b/letta/services/credit_verification_service.py
index cf4977e7..2efe1277 100644
--- a/letta/services/credit_verification_service.py
+++ b/letta/services/credit_verification_service.py
@@ -15,9 +15,13 @@ class CreditVerificationService:
         self.endpoint = os.getenv("STEP_ORCHESTRATOR_ENDPOINT")
         self.auth_key = os.getenv("STEP_COMPLETE_KEY")
 
-    async def verify_credits(self, organization_id: str) -> bool:
+    async def verify_credits(self, organization_id: str, agent_id: str) -> bool:
         """
-        Check if an organization has enough credits to proceed.
+        Check if an organization has enough credits to proceed with a specific agent.
+
+        Args:
+            organization_id: The organization's core ID
+            agent_id: The agent's ID (used to determine model-specific costs)
 
         Returns True if credits are available or if the service is not configured.
         Raises InsufficientCreditsError if no credits remain.
@@ -32,7 +36,7 @@ class CreditVerificationService:
 
             async with httpx.AsyncClient(timeout=5.0) as client:
                 response = await client.get(
-                    f"{self.endpoint}/validate/core-organizations/{organization_id}",
+                    f"{self.endpoint}/validate/core-organizations/{organization_id}/agents/{agent_id}",
                     headers=headers,
                 )
                 response.raise_for_status()
@@ -46,11 +50,11 @@ class CreditVerificationService:
         except InsufficientCreditsError:
             raise
         except httpx.TimeoutException:
-            logger.warning(f"Timeout verifying credits for organization {organization_id}")
+            logger.warning(f"Timeout verifying credits for organization {organization_id}, agent {agent_id}")
             return True
         except httpx.HTTPStatusError as e:
-            logger.warning(f"HTTP error verifying credits for organization {organization_id}: {e.response.status_code}")
+            logger.warning(f"HTTP error verifying credits for organization {organization_id}, agent {agent_id}: {e.response.status_code}")
             return True
         except Exception as e:
-            logger.error(f"Unexpected error verifying credits for organization {organization_id}: {e}")
+            logger.error(f"Unexpected error verifying credits for organization {organization_id}, agent {agent_id}: {e}")
             return True