feat(ask): return partial answers on budget exhaustion + resumable sessions

Instead of returning an error when the agent runs out of iterations,
nudge it to synthesize its best answer on the last iteration and return
that answer (with a session ID for continuation). Users can resume
budget-exhausted sessions via continue_from without re-spending tokens.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: leifericf

resources/prompts/agent-system.edn

@@ -44,4 +44,5 @@ These are pre-built named queries that demonstrate valid Datalog patterns. Use t
 5. If a query returns empty results or an error, reformulate and try again.
 6. When you have enough information, emit {:tool :answer :args {:text \"...\"}} with your answer.
 7. Be precise — cite specific entities, paths, or values from query results.
-8. Keep queries focused. Prefer multiple simple queries over one complex query. Chain results across queries to build up a complete answer."}
+8. Keep queries focused. Prefer multiple simple queries over one complex query. Chain results across queries to build up a complete answer.
+9. You have a limited iteration budget. Do not save your answer for later — if you have enough information to answer (even partially), emit {:tool :answer} promptly. You will be warned when you are about to run out of iterations."}

src/noumenon/agent.clj

@@ -184,10 +184,47 @@
       {:result (str "Unknown tool: " (:tool tool-call)
                     ". Available: :query, :schema, :rules, :answer")})))
 
+;; --- Session storage (for resumable budget-exhausted runs) ---
+
+(def ^:private max-sessions 20)
+(def ^:private session-ttl-ms (* 30 60 1000)) ;; 30 minutes
+
+(def ^:private sessions (atom {}))
+
+(defn- evict-expired-sessions! []
+  (let [cutoff (- (System/currentTimeMillis) session-ttl-ms)]
+    (swap! sessions (fn [m]
+                      (let [live (into {} (remove #(< (:created-at (val %)) cutoff)) m)]
+                        (if (<= (count live) max-sessions)
+                          live
+                          (->> live
+                               (sort-by (comp :created-at val))
+                               (drop (- (count live) max-sessions))
+                               (into {}))))))))
+
+(defn- store-session! [state]
+  (evict-expired-sessions!)
+  (let [id (str (java.util.UUID/randomUUID))]
+    (swap! sessions assoc id (assoc state :created-at (System/currentTimeMillis)))
+    id))
+
+(defn- load-session! [id]
+  (let [s (get @sessions id)]
+    (when s
+      (swap! sessions dissoc id)
+      (dissoc s :created-at))))
+
 ;; --- Agent loop ---
 
 (def ^:private default-max-iterations 10)
 
+(def ^:private budget-nudge
+  (str "You have reached your iteration budget. "
+       "You MUST respond with {:tool :answer :args {:text \"...\"}} NOW. "
+       "Synthesize the best answer you can from the information gathered so far. "
+       "If your answer is incomplete, end with a note about what additional queries "
+       "you would have run given more budget."))
+
 (defn- parse-error-transition
   [messages response-text parse-error]
   (let [error-msg (str "Your response could not be parsed as EDN. Error: "
@@ -203,16 +240,29 @@
         {:role "assistant" :content response-text}
         {:role "user" :content (str "Tool result:\n" result)}))
 
+(defn- maybe-append-nudge
+  "Append the budget nudge if this is the last allowed iteration."
+  [messages iterations max-iterations]
+  (if (>= iterations (dec max-iterations))
+    (conj messages {:role "user" :content budget-nudge})
+    messages))
+
 (defn- next-state
   [{:keys [db invoke-fn]}
    {:keys [messages steps iterations total-usage max-iterations]}]
   (if (>= iterations max-iterations)
-    {:done {:answer (or (some :answer steps)
-                        "Budget exhausted: reached maximum iterations without a final answer.")
-            :steps  steps
-            :usage  (assoc total-usage :iterations iterations)
-            :status :budget-exhausted}}
-    (let [response      (invoke-fn (vec messages))
+    (let [state {:messages messages :steps steps
+                 :iterations iterations :total-usage total-usage
+                 :max-iterations max-iterations}
+          session-id (store-session! state)]
+      {:done {:answer     (or (some :answer steps)
+                              "Budget exhausted: reached maximum iterations without a final answer.")
+              :steps      steps
+              :usage      (assoc total-usage :iterations iterations)
+              :status     :budget-exhausted
+              :session-id session-id}})
+    (let [msgs          (maybe-append-nudge messages iterations max-iterations)
+          response      (invoke-fn (vec msgs))
           usage         (merge-with + total-usage
                                     (select-keys (:usage response)
                                                  [:input-tokens :output-tokens :cost-usd :duration-ms]))
@@ -240,16 +290,20 @@
 
 (defn ask
   "Run the agent loop: prompt → parse → dispatch → repeat.
-   Returns {:answer string :steps vec :usage {:iterations n :input-tokens n :output-tokens n}}."
-  [db question {:keys [invoke-fn repo-name max-iterations]
+   Returns {:answer string :steps vec :usage {:iterations n :input-tokens n :output-tokens n}
+            :status :answered|:budget-exhausted :session-id string?}.
+   Pass :continue-from session-id to resume a budget-exhausted session."
+  [db question {:keys [invoke-fn repo-name max-iterations continue-from]
                 :or   {max-iterations default-max-iterations}}]
-  (let [system-prompt (build-system-prompt db repo-name)
-        context      {:db db :invoke-fn invoke-fn}
-        initial      {:messages [{:role "user" :content (str system-prompt "\n\n" question)}]
-                      :steps []
-                      :iterations 0
-                      :total-usage llm/zero-usage
-                      :max-iterations max-iterations}]
+  (let [context {:db db :invoke-fn invoke-fn}
+        initial (if-let [prev (when continue-from (load-session! continue-from))]
+                  (assoc prev :max-iterations (+ (:iterations prev) max-iterations))
+                  {:messages [{:role "user" :content (str (build-system-prompt db repo-name)
+                                                          "\n\n" question)}]
+                   :steps []
+                   :iterations 0
+                   :total-usage llm/zero-usage
+                   :max-iterations max-iterations})]
     (loop [state initial]
       (let [nxt (next-state context state)]
         (if-let [done (:done nxt)]

src/noumenon/cli.clj

@@ -217,6 +217,8 @@
                   :desc "Max query iterations (default: 10)"
                   :error-invalid :invalid-max-iterations
                   :error-missing :missing-max-iterations-value}
+                 {:flag "--continue-from" :key :continue-from :parse :string
+                  :desc "Session ID from a budget-exhausted run — resumes the agent"}
                  db-dir-flag]
                 verbose-flags))
    :initial {:subcommand "ask"}

src/noumenon/main.clj

@@ -237,7 +237,7 @@
 
 (defn do-ask
   "Run the ask subcommand. Returns {:exit n :result map-or-nil}."
-  [{:keys [question model provider max-iterations verbose] :as opts}]
+  [{:keys [question model provider max-iterations continue-from verbose] :as opts}]
   (with-valid-repo
     opts
     (fn [ctx]
@@ -252,7 +252,8 @@
                                                    :max-tokens  4096})
                   result (agent/ask db question
                                     (cond-> {:invoke-fn invoke-fn :repo-name db-name}
-                                      max-iterations (assoc :max-iterations max-iterations)))]
+                                      max-iterations (assoc :max-iterations max-iterations)
+                                      continue-from  (assoc :continue-from continue-from)))]
               (when verbose
                 (let [max-iters (or max-iterations 10)]
                   (doseq [step (:steps result)]
@@ -267,15 +268,27 @@
                                                            " chars)"))
                                 :else               "thinking")]
                       (log! (str "  [" i "/" max-iters "] " tag))))))
-              (let [exit-code (if (= :budget-exhausted (:status result)) 2 0)]
-                (if (= :budget-exhausted (:status result))
-                  (log! "Budget exhausted — no answer found.")
-                  (when-let [answer (:answer result)]
-                    (log! answer)))
-                {:exit   exit-code
-                 :result {:answer (:answer result)
-                          :status (:status result)
-                          :usage  (:usage result)}}))))
+              (let [exhausted? (= :budget-exhausted (:status result))
+                    answer     (:answer result)
+                    session-id (:session-id result)]
+                (cond
+                  (and exhausted? (not answer))
+                  (do (log! "Budget exhausted — no answer found.")
+                      {:exit 2 :result {:status :budget-exhausted :usage (:usage result)}})
+
+                  exhausted?
+                  (do (log! answer)
+                      (log! (str "\n[Session " session-id " saved — re-run with"
+                                 " --continue-from " session-id " to resume]"))
+                      {:exit   2
+                       :result {:answer answer :status :budget-exhausted
+                                :session-id session-id :usage (:usage result)}})
+
+                  :else
+                  (do (when answer (log! answer))
+                      {:exit   0
+                       :result {:answer answer :status (:status result)
+                                :usage (:usage result)}}))))))
         (catch clojure.lang.ExceptionInfo e
           (print-error! (.getMessage e))
           {:exit 1})))))

src/noumenon/mcp.clj

@@ -144,7 +144,8 @@
                                      {"question" {:type "string" :description "Question to ask about the repository"}
                                       "provider" {:type "string" :description "LLM provider: glm, claude-api, or claude-cli (aliases: claude = claude-cli)"}
                                       "model" {:type "string" :description "Model alias (e.g. sonnet, haiku, opus)"}
-                                      "max_iterations" {:type "integer" :description "Max query iterations (default: 10, max: 50)"}})
+                                      "max_iterations" {:type "integer" :description "Max query iterations (default: 10, max: 50)"}
+                                      "continue_from" {:type "string" :description "Session ID from a budget-exhausted run — resumes the agent from where it left off"}})
                   :required ["question" "repo_path"]}}
    {:name "noumenon_analyze"
     :description "Run LLM analysis on repository files to enrich the knowledge graph with semantic metadata. Only analyzes files not yet analyzed. Requires a prior import."
@@ -336,18 +337,25 @@
             result      (agent/ask db (args "question")
                                    {:invoke-fn      invoke-fn
                                     :repo-name      db-name
-                                    :max-iterations max-iter})
-            usage       (:usage result)]
+                                    :max-iterations max-iter
+                                    :continue-from  (args "continue_from")})
+            usage       (:usage result)
+            answer      (:answer result)
+            session-id  (:session-id result)]
         (log! "agent/done"
               (str "status=" (:status result)
                    " iterations=" (:iterations usage)
                    " tokens=" (+ (:input-tokens usage 0) (:output-tokens usage 0))))
         (if (= :budget-exhausted (:status result))
-          (tool-error (str "Budget exhausted after " max-iter " iterations"
-                           " (" (:input-tokens usage 0) " in / "
-                           (:output-tokens usage 0) " out tokens). "
-                           "Try increasing max_iterations or narrowing the question."))
-          (tool-result (or (:answer result)
+          (if answer
+            (tool-result (str answer
+                              "\n\n[Session " session-id " saved — to continue exploring, "
+                              "call noumenon_ask with continue_from=\"" session-id "\"]"))
+            (tool-error (str "Budget exhausted after " max-iter " iterations"
+                             " (" (:input-tokens usage 0) " in / "
+                             (:output-tokens usage 0) " out tokens) with no answer. "
+                             "Try increasing max_iterations or narrowing the question.")))
+          (tool-result (or answer
                            (str "No answer found (status: " (name (:status result)) ")"))))))))
 
 (defn- handle-analyze [args defaults]

test/noumenon/agent_test.clj

@@ -266,7 +266,56 @@
                              :max-iterations 3})]
     (is (= :budget-exhausted (:status result)))
     (is (<= @call-count 3))
-    (is (= 3 (get-in result [:usage :iterations])))))
+    (is (= 3 (get-in result [:usage :iterations])))
+    (is (string? (:session-id result)))))
+
+(deftest ask-sends-nudge-on-last-iteration
+  (testing "agent receives budget nudge message on final iteration"
+    (let [db            (make-test-db)
+          last-messages (atom nil)
+          call-count    (atom 0)
+          mock-llm      (fn [messages]
+                          (reset! last-messages messages)
+                          (swap! call-count inc)
+                          (if (>= @call-count 3)
+                            {:text  "{:tool :answer :args {:text \"partial answer\"}}"
+                             :usage {:input-tokens 100 :output-tokens 50}
+                             :model "mock"}
+                            {:text  "{:tool :query :args {:query [:find ?p :where [?e :file/path ?p]]}}"
+                             :usage {:input-tokens 100 :output-tokens 50}
+                             :model "mock"}))
+          result        (agent/ask db "test" {:invoke-fn mock-llm :repo-name "test"
+                                              :max-iterations 3})]
+      (is (= :answered (:status result)))
+      (is (= "partial answer" (:answer result)))
+      (is (re-find #"iteration budget" (:content (last @last-messages)))))))
+
+(deftest ask-session-store-and-continue
+  (testing "budget-exhausted session can be resumed with continue-from"
+    (let [db         (make-test-db)
+          call-count (atom 0)
+          mock-llm   (fn [_messages]
+                       (swap! call-count inc)
+                       (if (>= @call-count 4)
+                         {:text  "{:tool :answer :args {:text \"full answer\"}}"
+                          :usage {:input-tokens 100 :output-tokens 50}
+                          :model "mock"}
+                         {:text  "{:tool :query :args {:query [:find ?p :where [?e :file/path ?p]]}}"
+                          :usage {:input-tokens 100 :output-tokens 50}
+                          :model "mock"}))
+          ;; First run: exhaust budget after 2 iterations
+          result1    (agent/ask db "test" {:invoke-fn mock-llm :repo-name "test"
+                                           :max-iterations 2})
+          _          (is (= :budget-exhausted (:status result1)))
+          session-id (:session-id result1)
+          _          (is (some? session-id))
+          ;; Continue: resume with more budget
+          result2    (agent/ask db "test" {:invoke-fn      mock-llm
+                                           :repo-name      "test"
+                                           :max-iterations 5
+                                           :continue-from  session-id})]
+      (is (= :answered (:status result2)))
+      (is (= "full answer" (:answer result2))))))
 
 ;; --- Tier 1: integration test with mock LLM ---