From 417f56a5b66f94d0b8b40a2ace3731fb67c7ef79 Mon Sep 17 00:00:00 2001
From: Chuck Butkus <chuck@all-hands.dev>
Date: Wed, 18 Mar 2026 21:50:47 -0400
Subject: [PATCH] Fix when budgets are None

---
 enterprise/storage/lite_llm_manager.py | 60 +++++++++++++++++---------
 1 file changed, 39 insertions(+), 21 deletions(-)

diff --git a/enterprise/storage/lite_llm_manager.py b/enterprise/storage/lite_llm_manager.py
index 836ebe8278..dc38bb10cc 100644
--- a/enterprise/storage/lite_llm_manager.py
+++ b/enterprise/storage/lite_llm_manager.py
@@ -530,20 +530,26 @@ class LiteLlmManager:
         if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
             logger.warning('LiteLLM API configuration not found')
             return
+
+        json_data: dict[str, Any] = {
+            'team_id': team_id,
+            'team_alias': team_alias,
+            'models': [],
+            'spend': 0,
+            'metadata': {
+                'version': ORG_SETTINGS_VERSION,
+                'model': get_default_litellm_model(),
+            },
+        }
+
+        if max_budget is not None:
+            json_data['max_budget'] = max_budget
+
         response = await client.post(
             f'{LITE_LLM_API_URL}/team/new',
-            json={
-                'team_id': team_id,
-                'team_alias': team_alias,
-                'models': [],
-                'max_budget': max_budget,
-                'spend': 0,
-                'metadata': {
-                    'version': ORG_SETTINGS_VERSION,
-                    'model': get_default_litellm_model(),
-                },
-            },
+            json=json_data,
         )
+
         # Team failed to create in litellm - this is an unforseen error state...
         if not response.is_success:
             if (
@@ -923,14 +929,20 @@ class LiteLlmManager:
         if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
             logger.warning('LiteLLM API configuration not found')
             return
+
+        json_data: dict[str, Any] = {
+            'team_id': team_id,
+            'member': {'user_id': keycloak_user_id, 'role': 'user'},
+        }
+
+        if max_budget is not None:
+            json_data['max_budget_in_team'] = max_budget
+
         response = await client.post(
             f'{LITE_LLM_API_URL}/team/member_add',
-            json={
-                'team_id': team_id,
-                'member': {'user_id': keycloak_user_id, 'role': 'user'},
-                'max_budget_in_team': max_budget,
-            },
+            json=json_data,
         )
+
         # Failed to add user to team - this is an unforseen error state...
         if not response.is_success:
             if (
@@ -1003,14 +1015,20 @@ class LiteLlmManager:
         if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
             logger.warning('LiteLLM API configuration not found')
             return
+
+        json_data: dict[str, Any] = {
+            'team_id': team_id,
+            'user_id': keycloak_user_id,
+        }
+
+        if max_budget is not None:
+            json_data['max_budget_in_team'] = max_budget
+
         response = await client.post(
             f'{LITE_LLM_API_URL}/team/member_update',
-            json={
-                'team_id': team_id,
-                'user_id': keycloak_user_id,
-                'max_budget_in_team': max_budget,
-            },
+            json=json_data,
         )
+
         # Failed to update user in team - this is an unforseen error state...
         if not response.is_success:
             logger.error(