From b079fb03ad0214fe01f9b13b09c202f568d190bf Mon Sep 17 00:00:00 2001
From: Yetibot <yetibot@yetibot.com>
Date: Sun, 31 May 2026 14:46:39 +0000
Subject: [PATCH 1/2] feat(budget): incorporate veo budget into banana budget

---
 src/yetibot/core/commands/banana.clj          | 18 +++++++++++++++
 src/yetibot/core/commands/veo.clj             | 18 +++++++++++++++
 src/yetibot/core/util/gemini.clj              | 10 +++++++-
 .../core/test/commands/banana_test.clj        | 23 +++++++++++++++++++
 test/yetibot/core/test/commands/veo_test.clj  | 23 +++++++++++++++++++
 test/yetibot/core/test/util/gemini_test.clj   | 20 ++++++++++++++++
 6 files changed, 111 insertions(+), 1 deletion(-)
 create mode 100644 test/yetibot/core/test/commands/banana_test.clj
 create mode 100644 test/yetibot/core/test/commands/veo_test.clj
 create mode 100644 test/yetibot/core/test/util/gemini_test.clj

diff --git a/src/yetibot/core/commands/banana.clj b/src/yetibot/core/commands/banana.clj
index 91f05f60..1c6f706c 100644
--- a/src/yetibot/core/commands/banana.clj
+++ b/src/yetibot/core/commands/banana.clj
@@ -5,6 +5,23 @@
             [yetibot.core.util.gemini :as gemini]
             [yetibot.core.webapp.routes.images :refer [store-image!]]))
 
+(defn banana-budget-cmd
+  "banana budget # show monthly budget status"
+  {:yb/cat #{:info}}
+  [_]
+  (if (gemini/configured?)
+    (try
+      (let [{:keys [images-generated max-images spent budget remaining images-left veo-clips-left veo-cost-units month]}
+            (gemini/budget-status)]
+        {:result/value (format "Monthly Gemini budget status for %s:\n- Total Budget: $%.2f\n- Spent: $%.2f (%.1f%%)\n- Remaining: $%.2f\n- Image Units Generated: %d/%d\n- Remaining capacity: ~%d images OR ~%d Veo video clips (each clip costs %d image-units)"
+                               month budget spent (* 100 (/ spent budget)) remaining images-generated max-images images-left veo-clips-left veo-cost-units)
+         :result/data (gemini/budget-status)})
+      (catch Exception e
+        (error "banana budget error:" (.getMessage e))
+        {:result/error (str "Could not load budget status: " (.getMessage e))}))
+    {:result/error
+     "Gemini API is not configured. Set `gemini.key` in config."}))
+
 (defn banana-cmd
   "banana <prompt> # generate an image using Gemini nano banana image generation"
   {:yb/cat #{:img}}
@@ -29,4 +46,5 @@
      "Gemini API is not configured. Set `gemini.key` in config."}))
 
 (cmd-hook #"banana"
+  #"budget" banana-budget-cmd
   #".+" banana-cmd)
diff --git a/src/yetibot/core/commands/veo.clj b/src/yetibot/core/commands/veo.clj
index 22400d6e..5729987c 100644
--- a/src/yetibot/core/commands/veo.clj
+++ b/src/yetibot/core/commands/veo.clj
@@ -5,6 +5,23 @@
             [yetibot.core.util.gemini :as gemini]
             [yetibot.core.webapp.routes.images :refer [store-image!]]))
 
+(defn veo-budget-cmd
+  "veo budget # show monthly budget status"
+  {:yb/cat #{:info}}
+  [_]
+  (if (gemini/configured?)
+    (try
+      (let [{:keys [images-generated max-images spent budget remaining images-left veo-clips-left veo-cost-units month]}
+            (gemini/budget-status)]
+        {:result/value (format "Monthly Gemini budget status for %s:\n- Total Budget: $%.2f\n- Spent: $%.2f (%.1f%%)\n- Remaining: $%.2f\n- Image Units Generated: %d/%d\n- Remaining capacity: ~%d images OR ~%d Veo video clips (each clip costs %d image-units)"
+                               month budget spent (* 100 (/ spent budget)) remaining images-generated max-images images-left veo-clips-left veo-cost-units)
+         :result/data (gemini/budget-status)})
+      (catch Exception e
+        (error "veo budget error:" (.getMessage e))
+        {:result/error (str "Could not load budget status: " (.getMessage e))}))
+    {:result/error
+     "Gemini API is not configured. Set `gemini.key` in config."}))
+
 (defn veo-cmd
   "veo <prompt> # generate a short AI video with Veo
 
@@ -31,4 +48,5 @@
      "Gemini API is not configured. Set `gemini.key` in config."}))
 
 (cmd-hook #"veo"
+  #"budget" veo-budget-cmd
   #".+" veo-cmd)
diff --git a/src/yetibot/core/util/gemini.clj b/src/yetibot/core/util/gemini.clj
index b350fc05..fe086b59 100644
--- a/src/yetibot/core/util/gemini.clj
+++ b/src/yetibot/core/util/gemini.clj
@@ -18,6 +18,8 @@
 (s/def ::monthly (s/keys :opt-un [::budget]))
 (s/def ::budget (s/or :string string? :number number?))
 
+(declare veo-cost-units)
+
 (s/def ::config (s/keys :req-un [::key] :opt-un [::cost ::monthly]))
 
 ;; All Gemini settings live under [:gemini] and share a single API key
@@ -120,12 +122,18 @@
   (let [{:keys [count]} @usage-tracker
         max-imgs (max-images-per-month)
         spent (* count (cost-per-image))
-        remaining (- (monthly-budget) spent)]
+        remaining (max 0.0 (- (monthly-budget) spent))
+        v-units (veo-cost-units)
+        images-left (long (Math/floor (/ remaining (cost-per-image))))
+        veo-clips-left (long (Math/floor (/ images-left v-units)))]
     {:images-generated count
      :max-images max-imgs
      :spent (double spent)
      :budget (monthly-budget)
      :remaining (double remaining)
+     :images-left images-left
+     :veo-clips-left veo-clips-left
+     :veo-cost-units v-units
      :month (:month @usage-tracker)}))
 
 (defn- check-budget!
diff --git a/test/yetibot/core/test/commands/banana_test.clj b/test/yetibot/core/test/commands/banana_test.clj
new file mode 100644
index 00000000..d669d5f0
--- /dev/null
+++ b/test/yetibot/core/test/commands/banana_test.clj
@@ -0,0 +1,23 @@
+(ns yetibot.core.test.commands.banana-test
+  (:require [midje.sweet :refer [facts fact => contains provided]]
+            [yetibot.core.commands.banana :as b]
+            [yetibot.core.util.gemini :as gemini]))
+
+(facts "about banana-budget-cmd"
+       (fact "it returns an error if Gemini is not configured"
+             (b/banana-budget-cmd {}) => (contains {:result/error string?})
+             (provided (gemini/configured?) => false))
+
+       (fact "it returns budget status if configured"
+             (b/banana-budget-cmd {}) => (contains {:result/value string? :result/data map?})
+             (provided
+               (gemini/configured?) => true
+               (gemini/budget-status) => {:images-generated 0
+                                          :max-images 100
+                                          :spent 0.0
+                                          :budget 10.0
+                                          :remaining 10.0
+                                          :images-left 100
+                                          :veo-clips-left 20
+                                          :veo-cost-units 5
+                                          :month "2026-05"})))
diff --git a/test/yetibot/core/test/commands/veo_test.clj b/test/yetibot/core/test/commands/veo_test.clj
new file mode 100644
index 00000000..1e5ed2e2
--- /dev/null
+++ b/test/yetibot/core/test/commands/veo_test.clj
@@ -0,0 +1,23 @@
+(ns yetibot.core.test.commands.veo-test
+  (:require [midje.sweet :refer [facts fact => contains provided]]
+            [yetibot.core.commands.veo :as v]
+            [yetibot.core.util.gemini :as gemini]))
+
+(facts "about veo-budget-cmd"
+       (fact "it returns an error if Gemini is not configured"
+             (v/veo-budget-cmd {}) => (contains {:result/error string?})
+             (provided (gemini/configured?) => false))
+
+       (fact "it returns budget status if configured"
+             (v/veo-budget-cmd {}) => (contains {:result/value string? :result/data map?})
+             (provided
+               (gemini/configured?) => true
+               (gemini/budget-status) => {:images-generated 0
+                                          :max-images 100
+                                          :spent 0.0
+                                          :budget 10.0
+                                          :remaining 10.0
+                                          :images-left 100
+                                          :veo-clips-left 20
+                                          :veo-cost-units 5
+                                          :month "2026-05"})))
diff --git a/test/yetibot/core/test/util/gemini_test.clj b/test/yetibot/core/test/util/gemini_test.clj
new file mode 100644
index 00000000..9dfa2b5e
--- /dev/null
+++ b/test/yetibot/core/test/util/gemini_test.clj
@@ -0,0 +1,20 @@
+(ns yetibot.core.test.util.gemini-test
+  (:require
+   [midje.sweet :refer [=> fact facts contains provided anything]]
+   [yetibot.core.db.image-budget :as image-budget]
+   [yetibot.core.util.gemini :as gemini]))
+
+(facts "about budget-status"
+       (fact "it calculates budget status including veo details"
+             (let [status (gemini/budget-status)]
+               status => (contains {:images-generated integer?
+                                    :max-images integer?
+                                    :spent number?
+                                    :budget number?
+                                    :remaining number?
+                                    :images-left integer?
+                                    :veo-clips-left integer?
+                                    :veo-cost-units integer?
+                                    :month string?}))
+             (provided
+               (image-budget/query (contains {:where/map {:month anything}})) => [])))

From cbb2bc32ebd1994c582194acf7cbdae3e19a0071 Mon Sep 17 00:00:00 2001
From: Yetibot <yetibot@yetibot.com>
Date: Sun, 7 Jun 2026 19:25:33 +0000
Subject: [PATCH 2/2] fix(budget): fix tests and add Gemini agent prompt
 sessions budget

---
 src/yetibot/core/commands/agent.clj           |  5 +++-
 src/yetibot/core/commands/banana.clj          |  6 ++---
 src/yetibot/core/commands/veo.clj             |  6 ++---
 src/yetibot/core/util/gemini.clj              | 18 ++++++++++---
 .../core/test/commands/banana_test.clj        |  2 ++
 test/yetibot/core/test/commands/help.clj      |  9 ++++---
 test/yetibot/core/test/commands/veo_test.clj  |  2 ++
 test/yetibot/core/test/logging.clj            | 15 ++++++-----
 test/yetibot/core/test/util/gemini_test.clj   | 25 ++++++++++---------
 9 files changed, 56 insertions(+), 32 deletions(-)

diff --git a/src/yetibot/core/commands/agent.clj b/src/yetibot/core/commands/agent.clj
index a337d53d..155a42aa 100644
--- a/src/yetibot/core/commands/agent.clj
+++ b/src/yetibot/core/commands/agent.clj
@@ -23,7 +23,8 @@
    [yetibot.core.adapters.adapter :as a]
    [yetibot.core.chat :as chat]
    [yetibot.core.config :refer [get-config]]
-   [yetibot.core.hooks :refer [cmd-hook]])
+   [yetibot.core.hooks :refer [cmd-hook]]
+   [yetibot.core.util.gemini :as gemini])
   (:import
    [java.nio.file Files]
    [java.nio.file.attribute FileAttribute]
@@ -474,9 +475,11 @@
     (sweep-stale-workdirs! (agent-workdir-max-age-ms))
     (let [dir (work-dir target)]
       (try
+        (gemini/check-budget!)
         (let [context (when on-discord (thread-context context-channel))
               token (github-token)
               {:keys [response exit timed-out]} (run-gemini-agent dir request context mentions token)
+              _ (gemini/record-image-generated! (gemini/agent-cost-units))
               reply (cond
                       timed-out (say-timeout (quot (agent-timeout-ms) 60000))
                       (not (string/blank? response)) (say-final response (pr-urls response))
diff --git a/src/yetibot/core/commands/banana.clj b/src/yetibot/core/commands/banana.clj
index 1c6f706c..b2854962 100644
--- a/src/yetibot/core/commands/banana.clj
+++ b/src/yetibot/core/commands/banana.clj
@@ -11,10 +11,10 @@
   [_]
   (if (gemini/configured?)
     (try
-      (let [{:keys [images-generated max-images spent budget remaining images-left veo-clips-left veo-cost-units month]}
+      (let [{:keys [images-generated max-images spent budget remaining images-left veo-clips-left veo-cost-units agent-sessions-left agent-cost-units month]}
             (gemini/budget-status)]
-        {:result/value (format "Monthly Gemini budget status for %s:\n- Total Budget: $%.2f\n- Spent: $%.2f (%.1f%%)\n- Remaining: $%.2f\n- Image Units Generated: %d/%d\n- Remaining capacity: ~%d images OR ~%d Veo video clips (each clip costs %d image-units)"
-                               month budget spent (* 100 (/ spent budget)) remaining images-generated max-images images-left veo-clips-left veo-cost-units)
+        {:result/value (format "Monthly Gemini budget status for %s:\n- Total Budget: $%.2f\n- Spent: $%.2f (%.1f%%)\n- Remaining: $%.2f\n- Image Units Generated: %d/%d\n- Remaining capacity: ~%d images OR ~%d Veo video clips (each clip costs %d image-units) OR ~%d Agent prompt sessions (each session costs %d image-units)"
+                               month budget spent (* 100 (/ spent budget)) remaining images-generated max-images images-left veo-clips-left veo-cost-units agent-sessions-left agent-cost-units)
          :result/data (gemini/budget-status)})
       (catch Exception e
         (error "banana budget error:" (.getMessage e))
diff --git a/src/yetibot/core/commands/veo.clj b/src/yetibot/core/commands/veo.clj
index 5729987c..d2492e3e 100644
--- a/src/yetibot/core/commands/veo.clj
+++ b/src/yetibot/core/commands/veo.clj
@@ -11,10 +11,10 @@
   [_]
   (if (gemini/configured?)
     (try
-      (let [{:keys [images-generated max-images spent budget remaining images-left veo-clips-left veo-cost-units month]}
+      (let [{:keys [images-generated max-images spent budget remaining images-left veo-clips-left veo-cost-units agent-sessions-left agent-cost-units month]}
             (gemini/budget-status)]
-        {:result/value (format "Monthly Gemini budget status for %s:\n- Total Budget: $%.2f\n- Spent: $%.2f (%.1f%%)\n- Remaining: $%.2f\n- Image Units Generated: %d/%d\n- Remaining capacity: ~%d images OR ~%d Veo video clips (each clip costs %d image-units)"
-                               month budget spent (* 100 (/ spent budget)) remaining images-generated max-images images-left veo-clips-left veo-cost-units)
+        {:result/value (format "Monthly Gemini budget status for %s:\n- Total Budget: $%.2f\n- Spent: $%.2f (%.1f%%)\n- Remaining: $%.2f\n- Image Units Generated: %d/%d\n- Remaining capacity: ~%d images OR ~%d Veo video clips (each clip costs %d image-units) OR ~%d Agent prompt sessions (each session costs %d image-units)"
+                               month budget spent (* 100 (/ spent budget)) remaining images-generated max-images images-left veo-clips-left veo-cost-units agent-sessions-left agent-cost-units)
          :result/data (gemini/budget-status)})
       (catch Exception e
         (error "veo budget error:" (.getMessage e))
diff --git a/src/yetibot/core/util/gemini.clj b/src/yetibot/core/util/gemini.clj
index fe086b59..b212a902 100644
--- a/src/yetibot/core/util/gemini.clj
+++ b/src/yetibot/core/util/gemini.clj
@@ -68,6 +68,12 @@
 (defn- max-images-per-month []
   (long (Math/floor (/ (monthly-budget) (cost-per-image)))))
 
+(defn- agent-cost-per-session []
+  (or (parse-number (-> config :agent :cost-per-session)) 1.00))
+
+(defn agent-cost-units []
+  (max 1 (long (Math/ceil (/ (agent-cost-per-session) (cost-per-image))))))
+
 (defn- current-month []
   (.format (YearMonth/now) (DateTimeFormatter/ofPattern "yyyy-MM")))
 
@@ -124,8 +130,10 @@
         spent (* count (cost-per-image))
         remaining (max 0.0 (- (monthly-budget) spent))
         v-units (veo-cost-units)
+        a-units (agent-cost-units)
         images-left (long (Math/floor (/ remaining (cost-per-image))))
-        veo-clips-left (long (Math/floor (/ images-left v-units)))]
+        veo-clips-left (long (Math/floor (/ images-left v-units)))
+        agent-sessions-left (long (Math/floor (/ images-left a-units)))]
     {:images-generated count
      :max-images max-imgs
      :spent (double spent)
@@ -134,9 +142,11 @@
      :images-left images-left
      :veo-clips-left veo-clips-left
      :veo-cost-units v-units
+     :agent-sessions-left agent-sessions-left
+     :agent-cost-units a-units
      :month (:month @usage-tracker)}))
 
-(defn- check-budget!
+(defn check-budget!
   "Throw if the monthly image generation budget has been exhausted."
   []
   (reset-if-new-month!)
@@ -145,13 +155,13 @@
     (when (>= count max-imgs)
       (let [spent (* count (cost-per-image))]
         (throw (ex-info
-                (format "Monthly image budget exhausted: %d/%d images ($%.2f/$%.2f). Resets next month."
+                (format "Monthly Gemini budget exhausted: %d/%d image-units ($%.2f/$%.2f). Resets next month."
                         count max-imgs spent (monthly-budget))
                 {:type :budget-exceeded
                  :count count
                  :max max-imgs}))))))
 
-(defn- record-image-generated!
+(defn record-image-generated!
   "Record generated media against the monthly budget. Persists to db.
    `units` lets pricier generations (e.g. a Veo clip) count as multiple
    image-equivalents so the shared dollar budget stays accurate."
diff --git a/test/yetibot/core/test/commands/banana_test.clj b/test/yetibot/core/test/commands/banana_test.clj
index d669d5f0..c2d9092c 100644
--- a/test/yetibot/core/test/commands/banana_test.clj
+++ b/test/yetibot/core/test/commands/banana_test.clj
@@ -20,4 +20,6 @@
                                           :images-left 100
                                           :veo-clips-left 20
                                           :veo-cost-units 5
+                                          :agent-sessions-left 10
+                                          :agent-cost-units 10
                                           :month "2026-05"})))
diff --git a/test/yetibot/core/test/commands/help.clj b/test/yetibot/core/test/commands/help.clj
index 23ce8c8c..99d09a62 100644
--- a/test/yetibot/core/test/commands/help.clj
+++ b/test/yetibot/core/test/commands/help.clj
@@ -2,7 +2,8 @@
   (:require [midje.sweet :refer [facts fact => contains provided]]
             [yetibot.core.commands.help :as h]
             [yetibot.core.models.default-command :refer [fallback-enabled?
-                                                         fallback-help-text-override]]
+                                                         fallback-help-text-override
+                                                         configured-default-command]]
             [yetibot.core.models.help :refer [get-alias-docs get-docs get-docs-for
                                               fuzzy-get-docs-for]]))
 
@@ -11,7 +12,8 @@
  (fact
   "it will, by default, return text explaining how fallback commands are enabled
    and that the defalt command is `help`"
-  (h/fallback-help-text) => (contains "default command is `help`"))
+  (h/fallback-help-text) => (contains "default command is `help`")
+  (provided (configured-default-command) => "help"))
  
  (fact
   "it will tell you that fallback commands are disabled"
@@ -47,7 +49,8 @@
   (h/help-topics nil)
   => #"(?is)`help <command>`.*`category`.*default command is `help`.*available commands.*`one`.*`two`"
   (provided (get-docs) => {"two" 2
-                           "one" 1})))
+                           "one" 1}
+            (configured-default-command) => "help")))
 
 (facts
  "about help-for-topic"
diff --git a/test/yetibot/core/test/commands/veo_test.clj b/test/yetibot/core/test/commands/veo_test.clj
index 1e5ed2e2..c791e10f 100644
--- a/test/yetibot/core/test/commands/veo_test.clj
+++ b/test/yetibot/core/test/commands/veo_test.clj
@@ -20,4 +20,6 @@
                                           :images-left 100
                                           :veo-clips-left 20
                                           :veo-cost-units 5
+                                          :agent-sessions-left 10
+                                          :agent-cost-units 10
                                           :month "2026-05"})))
diff --git a/test/yetibot/core/test/logging.clj b/test/yetibot/core/test/logging.clj
index 341df7e9..c1ae5ba8 100644
--- a/test/yetibot/core/test/logging.clj
+++ b/test/yetibot/core/test/logging.clj
@@ -33,12 +33,15 @@
 (facts
  "about start"
  (fact
-  "it will default to level :info and rolling-appender enabled when
-   no configs are provided"
-  (let [config (l/start)]
-    config => (contains {:level :info})
-    (get-in config [:appenders :rolling-appender]) => (contains
-                                                       {:enabled? true}))))
+  "it will default to level :info when no configs are provided"
+  (l/start) => (contains {:level :info})
+  (provided
+    (get-config anything anything) => {:error true}))
+ (fact
+  "it will default to rolling-appender enabled when no configs are provided"
+  (get-in (l/start) [:appenders :rolling-appender]) => (contains {:enabled? true})
+  (provided
+    (get-config anything anything) => {:error true})))
 
 (facts
  "about log-path-config"
diff --git a/test/yetibot/core/test/util/gemini_test.clj b/test/yetibot/core/test/util/gemini_test.clj
index 9dfa2b5e..2e2e5358 100644
--- a/test/yetibot/core/test/util/gemini_test.clj
+++ b/test/yetibot/core/test/util/gemini_test.clj
@@ -5,16 +5,17 @@
    [yetibot.core.util.gemini :as gemini]))
 
 (facts "about budget-status"
-       (fact "it calculates budget status including veo details"
-             (let [status (gemini/budget-status)]
-               status => (contains {:images-generated integer?
-                                    :max-images integer?
-                                    :spent number?
-                                    :budget number?
-                                    :remaining number?
-                                    :images-left integer?
-                                    :veo-clips-left integer?
-                                    :veo-cost-units integer?
-                                    :month string?}))
+       (fact "it calculates budget status including veo and agent details"
+             (gemini/budget-status) => (contains {:images-generated integer?
+                                                  :max-images integer?
+                                                  :spent number?
+                                                  :budget number?
+                                                  :remaining number?
+                                                  :images-left integer?
+                                                  :veo-clips-left integer?
+                                                  :veo-cost-units integer?
+                                                  :agent-sessions-left integer?
+                                                  :agent-cost-units integer?
+                                                  :month string?})
              (provided
-               (image-budget/query (contains {:where/map {:month anything}})) => [])))
+               (image-budget/query anything) => [])))