PLTF-309: disable budget enforcement when ENABLE_BILLING=false (#13440)

Co-authored-by: openhands <openhands@all-hands.dev>

enterprise/storage/lite_llm_manager.py

@@ -29,14 +29,37 @@ KEY_VERIFICATION_TIMEOUT = 5.0
 # A very large number to represent "unlimited" until LiteLLM fixes their unlimited update bug.
 UNLIMITED_BUDGET_SETTING = 1000000000.0
 
-try:
+# Check if billing is enabled (defaults to false for enterprise deployments)
-    DEFAULT_INITIAL_BUDGET = float(os.environ.get('DEFAULT_INITIAL_BUDGET', 0.0))
+ENABLE_BILLING = os.environ.get('ENABLE_BILLING', 'false').lower() == 'true'
-    if DEFAULT_INITIAL_BUDGET < 0:
+
+
+def _get_default_initial_budget() -> float | None:
+    """Get the default initial budget for new teams.
+
+    When billing is disabled (ENABLE_BILLING=false), returns None to disable
+    budget enforcement in LiteLLM. When billing is enabled, returns the
+    DEFAULT_INITIAL_BUDGET environment variable value (default 0.0).
+
+    Returns:
+        float | None: The default budget, or None to disable budget enforcement.
+    """
+    if not ENABLE_BILLING:
+        return None
+
+    try:
+        budget = float(os.environ.get('DEFAULT_INITIAL_BUDGET', 0.0))
+        if budget < 0:
+            raise ValueError(
+                f'DEFAULT_INITIAL_BUDGET must be non-negative, got {budget}'
+            )
+        return budget
+    except ValueError as e:
         raise ValueError(
-            f'DEFAULT_INITIAL_BUDGET must be non-negative, got {DEFAULT_INITIAL_BUDGET}'
+            f'Invalid DEFAULT_INITIAL_BUDGET environment variable: {e}'
-        )
+        ) from e
-except ValueError as e:
+
-    raise ValueError(f'Invalid DEFAULT_INITIAL_BUDGET environment variable: {e}') from e
+
+DEFAULT_INITIAL_BUDGET: float | None = _get_default_initial_budget()
 
 
 def get_openhands_cloud_key_alias(keycloak_user_id: str, org_id: str) -> str:
@@ -110,12 +133,15 @@ class LiteLlmManager:
             ) as client:
                 # Check if team already exists and get its budget
                 # New users joining existing orgs should inherit the team's budget
-                team_budget: float = DEFAULT_INITIAL_BUDGET
+                # When billing is disabled, DEFAULT_INITIAL_BUDGET is None
+                team_budget: float | None = DEFAULT_INITIAL_BUDGET
                 try:
                     existing_team = await LiteLlmManager._get_team(client, org_id)
                     if existing_team:
                         team_info = existing_team.get('team_info', {})
-                        team_budget = team_info.get('max_budget', 0.0) or 0.0
+                        # Preserve None from existing team (no budget enforcement)
+                        existing_budget = team_info.get('max_budget')
+                        team_budget = existing_budget
                         logger.info(
                             'LiteLlmManager:create_entries:existing_team_budget',
                             extra={
@@ -525,8 +551,17 @@ class LiteLlmManager:
         client: httpx.AsyncClient,
         team_alias: str,
         team_id: str,
-        max_budget: float,
+        max_budget: float | None,
     ):
+        """Create a new team in LiteLLM.
+
+        Args:
+            client: The HTTP client to use.
+            team_alias: The alias for the team.
+            team_id: The ID for the team.
+            max_budget: The maximum budget for the team. When None, budget
+                enforcement is disabled (unlimited usage).
+        """
         if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
             logger.warning('LiteLLM API configuration not found')
             return
@@ -536,7 +571,7 @@ class LiteLlmManager:
                 'team_id': team_id,
                 'team_alias': team_alias,
                 'models': [],
-                'max_budget': max_budget,
+                'max_budget': max_budget,  # None disables budget enforcement
                 'spend': 0,
                 'metadata': {
                     'version': ORG_SETTINGS_VERSION,
@@ -918,8 +953,17 @@ class LiteLlmManager:
         client: httpx.AsyncClient,
         keycloak_user_id: str,
         team_id: str,
-        max_budget: float,
+        max_budget: float | None,
     ):
+        """Add a user to a team in LiteLLM.
+
+        Args:
+            client: The HTTP client to use.
+            keycloak_user_id: The user's Keycloak ID.
+            team_id: The team ID.
+            max_budget: The maximum budget for the user in the team. When None,
+                budget enforcement is disabled (unlimited usage).
+        """
         if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
             logger.warning('LiteLLM API configuration not found')
             return
@@ -928,7 +972,7 @@ class LiteLlmManager:
             json={
                 'team_id': team_id,
                 'member': {'user_id': keycloak_user_id, 'role': 'user'},
-                'max_budget_in_team': max_budget,
+                'max_budget_in_team': max_budget,  # None disables budget enforcement
             },
         )
         # Failed to add user to team - this is an unforseen error state...
@@ -998,8 +1042,17 @@ class LiteLlmManager:
         client: httpx.AsyncClient,
         keycloak_user_id: str,
         team_id: str,
-        max_budget: float,
+        max_budget: float | None,
     ):
+        """Update a user's budget in a team.
+
+        Args:
+            client: The HTTP client to use.
+            keycloak_user_id: The user's Keycloak ID.
+            team_id: The team ID.
+            max_budget: The maximum budget for the user in the team. When None,
+                budget enforcement is disabled (unlimited usage).
+        """
         if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
             logger.warning('LiteLLM API configuration not found')
             return
@@ -1008,7 +1061,7 @@ class LiteLlmManager:
             json={
                 'team_id': team_id,
                 'user_id': keycloak_user_id,
-                'max_budget_in_team': max_budget,
+                'max_budget_in_team': max_budget,  # None disables budget enforcement
             },
         )
         # Failed to update user in team - this is an unforseen error state...

enterprise/tests/unit/test_lite_llm_manager.py

@@ -38,8 +38,9 @@ class TestDefaultInitialBudget:
         if 'storage.lite_llm_manager' in sys.modules:
             del sys.modules['storage.lite_llm_manager']
 
-        # Clear the env var
+        # Clear the env vars
         os.environ.pop('DEFAULT_INITIAL_BUDGET', None)
+        os.environ.pop('ENABLE_BILLING', None)
 
         # Restore original module or reimport fresh
         if original_module is not None:
@@ -47,31 +48,56 @@ class TestDefaultInitialBudget:
         else:
             importlib.import_module('storage.lite_llm_manager')
 
-    def test_default_initial_budget_defaults_to_zero(self):
+    def test_default_initial_budget_none_when_billing_disabled(self):
-        """Test that DEFAULT_INITIAL_BUDGET defaults to 0.0 when env var not set."""
+        """Test that DEFAULT_INITIAL_BUDGET is None when billing is disabled."""
         # Temporarily remove the module so we can reimport with different env vars
         if 'storage.lite_llm_manager' in sys.modules:
             del sys.modules['storage.lite_llm_manager']
 
-        # Clear the env var and reimport
+        # Ensure billing is disabled (default) and reimport
+        os.environ.pop('ENABLE_BILLING', None)
+        os.environ.pop('DEFAULT_INITIAL_BUDGET', None)
+        module = importlib.import_module('storage.lite_llm_manager')
+        assert module.DEFAULT_INITIAL_BUDGET is None
+
+    def test_default_initial_budget_defaults_to_zero_when_billing_enabled(self):
+        """Test that DEFAULT_INITIAL_BUDGET defaults to 0.0 when billing is enabled."""
+        # Temporarily remove the module so we can reimport with different env vars
+        if 'storage.lite_llm_manager' in sys.modules:
+            del sys.modules['storage.lite_llm_manager']
+
+        # Enable billing and reimport
+        os.environ['ENABLE_BILLING'] = 'true'
         os.environ.pop('DEFAULT_INITIAL_BUDGET', None)
         module = importlib.import_module('storage.lite_llm_manager')
         assert module.DEFAULT_INITIAL_BUDGET == 0.0
 
-    def test_default_initial_budget_uses_env_var(self):
+    def test_default_initial_budget_uses_env_var_when_billing_enabled(self):
-        """Test that DEFAULT_INITIAL_BUDGET uses value from environment variable."""
+        """Test that DEFAULT_INITIAL_BUDGET uses value from environment variable when billing enabled."""
         if 'storage.lite_llm_manager' in sys.modules:
             del sys.modules['storage.lite_llm_manager']
 
+        os.environ['ENABLE_BILLING'] = 'true'
         os.environ['DEFAULT_INITIAL_BUDGET'] = '100.0'
         module = importlib.import_module('storage.lite_llm_manager')
         assert module.DEFAULT_INITIAL_BUDGET == 100.0
 
+    def test_default_initial_budget_ignores_env_var_when_billing_disabled(self):
+        """Test that DEFAULT_INITIAL_BUDGET returns None when billing disabled, ignoring env var."""
+        if 'storage.lite_llm_manager' in sys.modules:
+            del sys.modules['storage.lite_llm_manager']
+
+        os.environ.pop('ENABLE_BILLING', None)  # billing disabled by default
+        os.environ['DEFAULT_INITIAL_BUDGET'] = '100.0'
+        module = importlib.import_module('storage.lite_llm_manager')
+        assert module.DEFAULT_INITIAL_BUDGET is None
+
     def test_default_initial_budget_rejects_invalid_value(self):
         """Test that DEFAULT_INITIAL_BUDGET raises ValueError for invalid values."""
         if 'storage.lite_llm_manager' in sys.modules:
             del sys.modules['storage.lite_llm_manager']
 
+        os.environ['ENABLE_BILLING'] = 'true'
         os.environ['DEFAULT_INITIAL_BUDGET'] = 'abc'
         with pytest.raises(ValueError) as exc_info:
             importlib.import_module('storage.lite_llm_manager')
@@ -82,6 +108,7 @@ class TestDefaultInitialBudget:
         if 'storage.lite_llm_manager' in sys.modules:
             del sys.modules['storage.lite_llm_manager']
 
+        os.environ['ENABLE_BILLING'] = 'true'
         os.environ['DEFAULT_INITIAL_BUDGET'] = '-10.0'
         with pytest.raises(ValueError) as exc_info:
             importlib.import_module('storage.lite_llm_manager')