diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 0ac1547693e..77a6e021c81 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -210,7 +210,9 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): def _make_cache_key(self, session_id: str, agent_id: str) -> str: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - scope: Final = json.dumps((global_agent_registry.stable_agent_id(agent_id), session_id), separators=(",", ":")) + stable_agent_id: Final = global_agent_registry.stable_agent_id(agent_id) + canonical_agent_id: Final = stable_agent_id if isinstance(stable_agent_id, str) else agent_id + scope: Final = json.dumps((canonical_agent_id, session_id), separators=(",", ":")) return f"{{agent_session_budget:{scope}}}:spend" async def _get_current_spend(self, cache_key: str) -> float: diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index eb27f65fcc4..1b26c59fc83 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -184,7 +184,9 @@ class _PROXY_MaxIterationsHandler(CustomLogger): return f"{{session_iterations:{session_id}}}:count" from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - scope: Final = json.dumps((global_agent_registry.stable_agent_id(agent_id), session_id), separators=(",", ":")) + stable_agent_id: Final = global_agent_registry.stable_agent_id(agent_id) + canonical_agent_id: Final = stable_agent_id if isinstance(stable_agent_id, str) else agent_id + scope: Final = json.dumps((canonical_agent_id, session_id), separators=(",", ":")) return f"{{agent_session_iterations:{scope}}}:count" async def _increment_and_get(self, cache_key: str) -> int: