Actions are refused before they are sent if they would pass the daily budget

custom_components/jev/services.py

@@ -217,9 +217,25 @@ async def _ask(
         type(state).__name__,
         state,
     )
+    usage = entry.runtime_data.usage
+    usage.roll_over(dt_util.now().date())
     request_bytes = payload_bytes(state, questions, entry.runtime_data.model)
+    # The same check a context and the voice agent make: a script that loops on an
+    # action is the runaway the budget exists for.
+    estimate = usage.estimate_tokens(request_bytes)
+    if usage.would_exceed_with(estimate):
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="action_over_budget",
+            translation_placeholders={
+                "estimate": str(estimate),
+                "remaining": str(usage.remaining()),
+                "budget": str(usage.budget),
+            },
+        )
     try:
-        response = await entry.runtime_data.client.ask(state, questions)
+        with usage.reservation(estimate):
+            response = await entry.runtime_data.client.ask(state, questions)
     except JevAuthError as err:
         entry.async_start_reauth(hass)
         raise HomeAssistantError(
@@ -233,8 +249,6 @@ async def _ask(
             translation_key="ask_failed",
             translation_placeholders={"reason": str(err)},
         ) from err
-    usage = entry.runtime_data.usage
-    usage.roll_over(dt_util.now().date())
     usage.record(response.usage.input_tokens, request_bytes)
     entry.runtime_data.model_version = response.model or entry.runtime_data.model_version
     usage.notify()

custom_components/jev/strings.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Task {task} costs about {estimate} input tokens and {remaining} of the {budget} daily budget are left. Raise or clear the budget in the integration options to continue."
     },
+    "action_over_budget": {
+      "message": "This action costs about {estimate} input tokens and {remaining} of the {budget} daily budget are left. Raise or clear the budget in the integration options to continue."
+    },
     "reserved_field": {
       "message": "The field {field} cannot be used, because the result carries the confidence behind every answer under that name. Rename the field."
     },

custom_components/jev/translations/cs.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Úloha {task} stojí asi {estimate} vstupních tokenů a z denního rozpočtu {budget} zbývá {remaining}. Pro pokračování zvyšte nebo zrušte rozpočet v možnostech integrace."
     },
+    "action_over_budget": {
+      "message": "Tato akce stojí asi {estimate} vstupních tokenů a z denního rozpočtu {budget} zbývá {remaining}. Pokud chcete pokračovat, zvyšte nebo zrušte rozpočet v možnostech integrace."
+    },
     "reserved_field": {
       "message": "Pole {field} nelze použít, protože výsledek pod tímto názvem nese jistotu každé odpovědi. Pole přejmenujte."
     },

custom_components/jev/translations/da.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Opgaven {task} koster omkring {estimate} inputtokens, og der er {remaining} tilbage af dagsbudgettet på {budget}. Hæv eller ryd budgettet i integrationens indstillinger for at fortsætte."
     },
+    "action_over_budget": {
+      "message": "Denne handling koster omkring {estimate} input-tokens, og der er {remaining} tilbage af det daglige budget på {budget}. Hæv eller fjern budgettet i integrationens indstillinger for at fortsætte."
+    },
     "reserved_field": {
       "message": "Feltet {field} kan ikke bruges, fordi resultatet bærer sikkerheden bag hvert svar under det navn. Omdøb feltet."
     },

custom_components/jev/translations/de.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Die Aufgabe {task} kostet etwa {estimate} Eingabetokens, und vom Tagesbudget von {budget} sind noch {remaining} übrig. Erhöhe oder entferne das Budget in den Optionen der Integration, um fortzufahren."
     },
+    "action_over_budget": {
+      "message": "Diese Aktion kostet etwa {estimate} Eingabe-Tokens, und vom Tagesbudget von {budget} sind noch {remaining} übrig. Erhöhe oder entferne das Budget in den Optionen der Integration, um fortzufahren."
+    },
     "reserved_field": {
       "message": "Das Feld {field} kann nicht verwendet werden, weil das Ergebnis unter diesem Namen die Sicherheit hinter jeder Antwort mitliefert. Benenne das Feld um."
     },

custom_components/jev/translations/en.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Task {task} costs about {estimate} input tokens and {remaining} of the {budget} daily budget are left. Raise or clear the budget in the integration options to continue."
     },
+    "action_over_budget": {
+      "message": "This action costs about {estimate} input tokens and {remaining} of the {budget} daily budget are left. Raise or clear the budget in the integration options to continue."
+    },
     "reserved_field": {
       "message": "The field {field} cannot be used, because the result carries the confidence behind every answer under that name. Rename the field."
     },

custom_components/jev/translations/es.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "La tarea {task} cuesta unos {estimate} tokens de entrada y quedan {remaining} del presupuesto diario de {budget}. Aumenta o borra el presupuesto en las opciones de la integración para continuar."
     },
+    "action_over_budget": {
+      "message": "Esta acción cuesta unos {estimate} tokens de entrada y quedan {remaining} del presupuesto diario de {budget}. Aumenta o elimina el presupuesto en las opciones de la integración para continuar."
+    },
     "reserved_field": {
       "message": "El campo {field} no se puede usar, porque el resultado lleva bajo ese nombre la confianza detrás de cada respuesta. Cambia el nombre del campo."
     },

custom_components/jev/translations/fr.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "La tâche {task} coûte environ {estimate} jetons d'entrée et il reste {remaining} sur le budget quotidien de {budget}. Augmentez ou supprimez le budget dans les options de l'intégration pour continuer."
     },
+    "action_over_budget": {
+      "message": "Cette action coûte environ {estimate} jetons d'entrée et il reste {remaining} sur le budget quotidien de {budget}. Augmentez ou supprimez le budget dans les options de l'intégration pour continuer."
+    },
     "reserved_field": {
       "message": "Le champ {field} ne peut pas être utilisé, car le résultat transporte sous ce nom la confiance derrière chaque réponse. Renommez le champ."
     },

custom_components/jev/translations/it.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "L'attività {task} costa circa {estimate} token di input e restano {remaining} del budget giornaliero di {budget}. Aumenta o cancella il budget nelle opzioni dell'integrazione per continuare."
     },
+    "action_over_budget": {
+      "message": "Questa azione costa circa {estimate} token di input e restano {remaining} del budget giornaliero di {budget}. Aumenta o rimuovi il budget nelle opzioni dell'integrazione per continuare."
+    },
     "reserved_field": {
       "message": "Il campo {field} non può essere usato, perché il risultato porta con quel nome la confidenza dietro ogni risposta. Rinomina il campo."
     },

custom_components/jev/translations/nl.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Taak {task} kost ongeveer {estimate} invoertokens en van het dagbudget van {budget} zijn er nog {remaining} over. Verhoog of wis het budget in de opties van de integratie om door te gaan."
     },
+    "action_over_budget": {
+      "message": "Deze actie kost ongeveer {estimate} invoertokens en er is nog {remaining} over van het dagbudget van {budget}. Verhoog of wis het budget in de opties van de integratie om door te gaan."
+    },
     "reserved_field": {
       "message": "Het veld {field} kan niet worden gebruikt, omdat het resultaat onder die naam de zekerheid achter elk antwoord meegeeft. Geef het veld een andere naam."
     },

custom_components/jev/translations/pl.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Zadanie {task} kosztuje około {estimate} tokenów wejściowych, a z dziennego budżetu {budget} zostało {remaining}. Zwiększ lub wyczyść budżet w opcjach integracji, aby kontynuować."
     },
+    "action_over_budget": {
+      "message": "Ta akcja kosztuje około {estimate} tokenów wejściowych, a z dziennego budżetu {budget} zostało {remaining}. Aby kontynuować, zwiększ lub usuń budżet w opcjach integracji."
+    },
     "reserved_field": {
       "message": "Pola {field} nie można użyć, ponieważ wynik przenosi pod tą nazwą pewność każdej odpowiedzi. Zmień nazwę pola."
     },

custom_components/jev/translations/pt-BR.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "A tarefa {task} custa cerca de {estimate} tokens de entrada e restam {remaining} do orçamento diário de {budget}. Aumente ou limpe o orçamento nas opções da integração para continuar."
     },
+    "action_over_budget": {
+      "message": "Esta ação custa cerca de {estimate} tokens de entrada e restam {remaining} do orçamento diário de {budget}. Aumente ou remova o orçamento nas opções da integração para continuar."
+    },
     "reserved_field": {
       "message": "O campo {field} não pode ser usado, porque o resultado leva sob esse nome a confiança por trás de cada resposta. Renomeie o campo."
     },

custom_components/jev/translations/ru.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Задача {task} стоит около {estimate} входных токенов, а из дневного бюджета {budget} осталось {remaining}. Чтобы продолжить, увеличьте или очистите бюджет в параметрах интеграции."
     },
+    "action_over_budget": {
+      "message": "Это действие стоит около {estimate} входных токенов, а от дневного бюджета {budget} осталось {remaining}. Чтобы продолжить, увеличьте или снимите бюджет в параметрах интеграции."
+    },
     "reserved_field": {
       "message": "Поле {field} использовать нельзя, потому что под этим именем результат передаёт уверенность каждого ответа. Переименуйте поле."
     },

custom_components/jev/translations/sv.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "Uppgiften {task} kostar ungefär {estimate} indatatokens och {remaining} av dagsbudgeten på {budget} återstår. Höj eller rensa budgeten i integrationens alternativ för att fortsätta."
     },
+    "action_over_budget": {
+      "message": "Den här åtgärden kostar ungefär {estimate} indatatokens och det finns {remaining} kvar av den dagliga budgeten på {budget}. Höj eller ta bort budgeten i integrationens alternativ för att fortsätta."
+    },
     "reserved_field": {
       "message": "Fältet {field} kan inte användas, eftersom resultatet bär säkerheten bakom varje svar under det namnet. Byt namn på fältet."
     },

custom_components/jev/translations/zh-Hans.json

@@ -306,6 +306,9 @@
     "task_over_budget": {
       "message": "任务 {task} 大约需要 {estimate} 个输入令牌，而每日预算 {budget} 只剩 {remaining}。请在集成选项中提高或清除预算后再继续。"
     },
+    "action_over_budget": {
+      "message": "此操作约需 {estimate} 个输入令牌，每日预算 {budget} 还剩 {remaining}。要继续，请在集成选项中提高或清除预算。"
+    },
     "reserved_field": {
       "message": "字段 {field} 不能使用，因为结果在这个名称下携带每个答案背后的置信度。请重命名该字段。"
     },

tests/test_services.py

@@ -437,3 +437,39 @@ async def test_an_action_with_two_entries_and_none_named_is_refused(
         {"state": "x", "instructions": "y", "config_entry": second.entry_id},
     )
     assert "noul" in named
+
+
+async def test_an_action_that_would_pass_the_budget_is_not_sent(
+    hass, loaded_entry, mock_client
+):
+    usage = loaded_entry.runtime_data.usage
+    usage.budget = 1000
+    usage.input_tokens = 999
+    mock_client.ask.reset_mock()
+
+    with pytest.raises(HomeAssistantError) as err:
+        await call(hass, "noul", {"state": "x", "instructions": "Is it done?"})
+
+    assert mock_client.ask.await_count == 0
+    assert err.value.translation_key == "action_over_budget"
+    placeholders = err.value.translation_placeholders
+    assert placeholders["remaining"] == "1"
+    assert placeholders["budget"] == "1000"
+    assert int(placeholders["estimate"]) > 1
+
+
+async def test_an_action_holds_its_estimate_while_it_runs(
+    hass, loaded_entry, mock_client
+):
+    usage = loaded_entry.runtime_data.usage
+    held = []
+
+    async def watch(state, questions):
+        held.append(usage.reserved)
+        return build_response(answer=NoulAnswer(noul=0.5))
+
+    mock_client.ask.side_effect = watch
+    await call(hass, "noul", {"state": "x", "instructions": "Is it done?"})
+
+    assert held[0] > 0
+    assert usage.reserved == 0