add_budget(): explicitly extend an exhausted session's hard budget

The mechanical caps bind follow-up tell() calls too - by design - so a
finished session can refuse more work. add_budget(steps, seconds) is
the explicit remedy; the refusal message now tells the agent to point
the user at it, and the extension persists to meta for resumes.

Author: mattyoreilly

R/session.R

@@ -307,6 +307,38 @@ atlas_session <- R6::R6Class("atlas_session",
       )
     },
 
+    #' @description Grant the agent more mechanical budget. The hard caps
+    #'   (`max_steps`, `max_runtime`) protect unattended runs, but they also
+    #'   bind follow-up `$tell()` calls on a finished session - top the
+    #'   budget up explicitly when you want more work done:
+    #'   `res$session$add_budget(steps = 25)`.
+    #' @param steps Additional code executions to allow.
+    #' @param seconds Additional wall-clock seconds to allow.
+    add_budget = function(steps = 0, seconds = 0) {
+      stopifnot(is.numeric(steps), steps >= 0,
+                is.numeric(seconds), seconds >= 0)
+      private$max_steps <- private$max_steps + steps
+      if (seconds > 0 && !is.null(private$deadline)) {
+        private$deadline <- max(private$deadline, Sys.time()) + seconds
+        private$max_runtime <- private$max_runtime + seconds
+      }
+      private$meta$max_steps <- private$max_steps
+      private$meta$max_runtime <- private$max_runtime
+      saveRDS(private$meta, file.path(self$dir, "meta.rds"))
+      if (private$verbose) {
+        cli::cli_alert_info(sprintf(
+          "Budget extended: %s steps remaining%s.",
+          if (is.finite(private$max_steps))
+            format(private$max_steps - private$steps) else "unlimited",
+          if (!is.null(private$deadline))
+            sprintf(", %.0f minutes on the clock",
+                    as.numeric(difftime(private$deadline, Sys.time(),
+                                        units = "mins")))
+          else ""))
+      }
+      invisible(self)
+    },
+
     #' @description Compact the conversation to save tokens: archive the
     #'   transcript to the run directory, clear the context window, and
     #'   re-orient the agent with a state briefing on the next message. The
@@ -523,7 +555,9 @@ atlas_session <- R6::R6Class("atlas_session",
               "blocked - this is enforced mechanically, do not retry. ",
               "Objects already in the session (atlas_models, ",
               "atlas_leaderboard) remain in place. Write your final report ",
-              "now from what you already know."))
+              "now from what you already know, and tell the user they can ",
+              "grant more budget with session$add_budget(steps = ...) if ",
+              "they want this work completed."))
           }
           private$steps <- private$steps + 1
           code <- trimws(code)

example.R

@@ -11,8 +11,8 @@
 library(atlas)
 
 # 1. Load your data ------------------------------------------------------
-csv_path <- "path/to/your_data.csv"   # <- edit me
-outcome  <- "DEFAULTED"               # <- edit me: the column to predict
+csv_path <- "~/Desktop/Hastings/HastingsCredit/Data/app_approved_train.csv"
+outcome <- "DEFAULTED" # <- edit me: the column to predict
 
 data <- read.csv(csv_path, stringsAsFactors = TRUE)
 str(data, list.len = 25)
@@ -24,29 +24,31 @@ atlas_leakage_screen(data, outcome)
 # Columns that won't exist at prediction time in deployment (IDs,
 # post-outcome fields, decline reasons...) - they are removed before the
 # agent ever sees the data:
-exclude <- c()                        # <- e.g. c("QUOTE_DECLINE_REASON")
+exclude <- c("default_rate", "Total_Loan_Amount", "PREDICTED_LOSS", "APR") # <- e.g. c("QUOTE_DECLINE_REASON")
 
 # 3. Interactive build ----------------------------------------------------
 # The agent explores, proposes a plan, and STOPS IN THE CONSOLE for your
 # approval - answer "yes", or steer it ("only GLMs, no trees").
 # Watch for the live tally lines as models and tweaks are scored:
 #   [tally #4 | gbm1: auc = 0.81 | best: glm2 = 0.79 | flat: 0/3 -> KEEP]
 res <- atlas(
-  data, outcome,
-  n_models  = 3,                      # maximum candidates
-  goal      = "prioritise interpretability",
-  exclude   = exclude,
-  test_prop = 0.2,                    # held out; agent never sees these rows
-  max_steps = 120                     # hard safety cap on code executions
+  data,
+  outcome,
+  n_models = 1, # maximum candidates
+  goal = "prioritise interpretability",
+  exclude = exclude,
+  test_prop = 0.2, # held out; agent never sees these rows
+  max_steps = 10, # hard safety cap on code executions
+  dir = "~/Desktop/Code"
 )
 
 # 4. What you got back ----------------------------------------------------
-res                    # leaderboards, tally summary, constraints, report, cost
-res$tally              # every attempt: KEEP / DISCARD, best-so-far
-res$test_leaderboard   # final ranking on the held-out rows (the one to trust)
-res$models             # fitted models: predict(res$models[[1]], newdata)
-res$dir                # run directory: report.md, code.R, tally.csv, plots
-cat(res$code, sep = "\n\n")   # the full script the agent wrote
+res # leaderboards, tally summary, constraints, report, cost
+res$tally # every attempt: KEEP / DISCARD, best-so-far
+res$test_leaderboard # final ranking on the held-out rows (the one to trust)
+res$models # fitted models: predict(res$models[[1]], newdata)
+res$dir # run directory: report.md, code.R, tally.csv, plots
+cat(res$code, sep = "\n\n") # the full script the agent wrote
 
 # The session is still live - ask it anything:
 res$session$tell("why did the winning model win?")

man/atlas_session.Rd

@@ -216,6 +216,22 @@ test_that("max_steps is a mechanical stop, not a suggestion", {
   expect_no_match(new_session()$chat$get_system_prompt(), "Hard budget")
 })
 
+test_that("add_budget() unblocks an exhausted session", {
+  dir <- temp_dir()
+  s <- atlas_session$new(mtcars, "mpg", chat = real_chat(), dir = dir,
+                         max_steps = 1)
+  run_tool <- s$chat$get_tools()$run_r_code
+  run_tool("1 + 1")
+  blocked <- run_tool("2 + 2")
+  expect_match(blocked, "BUDGET EXHAUSTED")
+  expect_match(blocked, "add_budget")   # the agent relays the remedy
+
+  s$add_budget(steps = 2)
+  expect_match(run_tool("2 + 2"), "^\\[1\\] 4")
+  # the extension is persisted for resumes
+  expect_equal(readRDS(file.path(dir, "meta.rds"))$max_steps, 3)
+})
+
 test_that("max_runtime blocks execution after the deadline", {
   s <- new_session(max_runtime = 0.01)     # expires almost immediately
   Sys.sleep(0.05)

tests/testthat/test-session.R

@@ -81,6 +81,16 @@ gone.
 Use them together: generous statistical rules so the agent explores, and a
 mechanical ceiling so "overnight" can't become "over the weekend".
 
+The budget binds the whole session, follow-up `$tell()` calls included - a
+finished run with nothing left in the tank will refuse further work. That
+is deliberate (a cap you can talk your way past is not a cap), and the
+remedy is explicit: grant more, then continue.
+
+```{r}
+res$session$add_budget(steps = 25)
+res$session$tell("remove the worst predictor and re-evaluate")
+```
+
 ## The live tally
 
 Every model and every tweak is scored the moment it is evaluated. The agent