feat: record token usage, cost and latency per agent and per run

demos/e2e-tests/e2e-tests.nlogox

@@ -276,6 +276,33 @@ to write-throttled-config [src dst cap interval-ms]
   file-close
 end
 
+;; ---------------------------------------------------------------------------
+;; T9 — token usage (#49)
+;;
+;; Every provider reports how many tokens a call consumed; the extension used
+;; to discard it. One real call must show up in the calling agent's usage
+;; and in the run total, with non-zero counts and a measured latency.
+;; ---------------------------------------------------------------------------
+
+to test-usage
+  log-line "T9 token usage"
+  let calls-before llm:get llm:usage "calls"
+  let run-before llm:get llm:usage-total "total-tokens"
+  carefully
+    [ let ignored llm:chat "Reply with exactly the word: OK" ]
+    [ log-line (word "  chat failed: " error-message) ]
+  let u llm:usage
+  log-line (word "  -> " u)
+  assert "one more call was recorded" (llm:get u "calls" = calls-before + 1)
+  assert "input tokens were reported" (llm:get u "input-tokens" > 0)
+  assert "output tokens were reported" (llm:get u "output-tokens" > 0)
+  assert "total equals input plus output"
+         (llm:get u "total-tokens" = (llm:get u "input-tokens") + (llm:get u "output-tokens"))
+  assert "latency was measured" (llm:get u "latency-ms" > 0)
+  assert "run total grew by this call" (llm:get llm:usage-total "total-tokens" > run-before)
+  report-totals
+end
+
 ;; Run everything in sequence.
 to run-headless
   carefully [ file-delete "e2e-results.txt" ] [ ]
@@ -298,6 +325,7 @@ to test-all
   test-structured-output
   test-choose
   test-throttling
+  test-usage
   output-print ""
   log-line (word "TOTAL  passed " pass-count "   failed " fail-count)
 end

docs/API-REFERENCE.md

@@ -20,6 +20,8 @@ The NetLogo Multi-LLM Extension provides a unified interface for multiple Large
 | `llm:set-thinking bool`            | Reasoning     | Enable/disable reasoning mode for current provider         |
 | `llm:set-reasoning-effort level`   | Reasoning     | Set effort: `"low"`, `"medium"`, `"high"`                  |
 | `llm:set-thinking-budget n`        | Reasoning     | Token budget for thinking (min 1024; Anthropic + Gemini)   |
+| `llm:usage`                        | Usage         | Token counts, cost, latency and call count for this agent  |
+| `llm:usage-total`                  | Usage         | The same totals across every agent in the run              |
 | `llm:history`                      | History       | Get current agent's conversation history                   |
 | `llm:set-history list`             | History       | Set conversation history for current agent                 |
 | `llm:clear-history`                | History       | Clear conversation history for current agent               |
@@ -858,6 +860,75 @@ print llm:list-models  ; Shows all providers, with Anthropic marked as ACTIVE
 - Custom models added via `models-override.yaml` are marked with `[custom]`
 - The currently active provider and model are marked with `[ACTIVE]`
 
+## Token Usage
+
+Every provider reports how many tokens a call consumed. The extension records
+that for each call and sums it per agent and for the whole run. Nothing is
+estimated: counts come from the provider's reply, cost only appears when a
+provider reports one, and latency is measured by the extension.
+
+### llm:usage
+
+**Syntax**: `llm:usage`
+
+**Description**: Token accounting for the calling agent since the last `clear-all`
+
+**Returns**: List of `[key value]` pairs, readable with `llm:get`
+
+| Key                  | Meaning                                                                 |
+| -------------------- | ----------------------------------------------------------------------- |
+| `input-tokens`       | Every prompt token the provider processed, cached tokens included       |
+| `output-tokens`      | Every generated token, reasoning tokens included                        |
+| `total-tokens`       | The provider's total, or input plus output when it reports none         |
+| `reasoning-tokens`   | Tokens spent on reasoning where reported separately, else 0             |
+| `cache-read-tokens`  | Prompt tokens served from a provider cache where reported, else 0       |
+| `cache-write-tokens` | Prompt tokens written to a provider cache where reported, else 0        |
+| `cost`               | Dollars, only when the provider reports it (OpenRouter); otherwise `""` |
+| `latency-ms`         | Summed wall time of the calls, including throttle waits and retries     |
+| `calls`              | How many provider replies contributed to these totals                   |
+
+**Example**:
+
+```netlogo
+ask turtles [
+  let reply llm:chat "Should I forage or rest?"
+  if llm:get llm:usage "total-tokens" > 5000 [
+    set color red   ;; this agent is expensive
+  ]
+]
+print (word "Run so far: " llm:get llm:usage-total "total-tokens" " tokens over "
+            llm:get llm:usage-total "calls" " calls")
+```
+
+**Notes**:
+
+- A reply the extension then rejects, such as a schema reply that does not
+  parse or an `llm:choose` answer that matches no option, still counts. The
+  provider billed those tokens. A call that never produced a reply records
+  nothing.
+- `cost` is `""` rather than `0` when unknown so an unknown cost is never
+  mistaken for a free call. Guard with `is-number?` before arithmetic.
+- Counters are cumulative. To measure one call, read `llm:usage` before and
+  after and subtract.
+- Providers define their fields differently; the mapping is normalised so the
+  headline numbers compare across providers:
+
+| Provider          | input-tokens                                             | output-tokens                          |
+| ----------------- | -------------------------------------------------------- | -------------------------------------- |
+| OpenAI-compatible | `prompt_tokens` (cached tokens are a subset)             | `completion_tokens` (reasoning subset) |
+| Anthropic         | `input_tokens` + cache creation + cache read             | `output_tokens` (thinking included)    |
+| Gemini            | `promptTokenCount` (cached content included)             | `candidatesTokenCount` + `thoughtsTokenCount` |
+| Ollama            | `prompt_eval_count`                                      | `eval_count`                           |
+
+### llm:usage-total
+
+**Syntax**: `llm:usage-total`
+
+**Description**: The same accounting summed over every agent in the run since
+the last `clear-all`, including agents that have since died
+
+**Returns**: List of `[key value]` pairs with the keys listed under `llm:usage`
+
 ## Structured Output Details
 
 ### Why JSON becomes nested lists

src/main/LLMExtension.scala

@@ -4,7 +4,7 @@ import org.nlogo.api._
 import org.nlogo.core.{LogoList, Syntax}
 import org.nlogo.extensions.llm.config.{ConfigLoader, ConfigStore}
 import org.nlogo.extensions.llm.providers.{LLMProvider, ProviderDescriptor, ProviderFactory, ProviderRegistry, ProviderRegistrations, ModelRegistry, OllamaProvider, ReadinessCheck, RetryPolicy}
-import org.nlogo.extensions.llm.models.{ChatMessage, ChatResponse, EnumFormat, JsonObjectFormat, ResponseFormat}
+import org.nlogo.extensions.llm.models.{ChatMessage, ChatResponse, EnumFormat, JsonObjectFormat, ResponseFormat, Usage}
 import org.nlogo.extensions.llm.utils.JsonToNetLogo
 import scala.collection.mutable.{ArrayBuffer, WeakHashMap}
 import scala.concurrent.{Await, ExecutionContext, Future}
@@ -61,6 +61,13 @@ class LLMExtension extends DefaultClassManager {
   // or provider calls; critical sections are small snapshots/appends only.
   private val historyLock = new Object
 
+  // Token accounting. Per-agent totals live in a WeakHashMap like history so a
+  // dead turtle's counters are collected with it; the run total is kept
+  // separately so it survives agent death. Both reset on clear-all.
+  private val agentUsage: WeakHashMap[Agent, UsageTotals] = WeakHashMap()
+  private var runUsage: UsageTotals = UsageTotals.empty
+  private val usageLock = new Object
+
   // Execution context for async operations
   implicit private val ec: ExecutionContext = ExecutionContext.global
 
@@ -129,13 +136,21 @@ class LLMExtension extends DefaultClassManager {
     manager.addPrimitive("list-models", ListModelsReporter)
     manager.addPrimitive("active", ActiveReporter)
     manager.addPrimitive("config", ConfigReporter)
+
+    // Token accounting primitives
+    manager.addPrimitive("usage", UsageReporter)
+    manager.addPrimitive("usage-total", UsageTotalReporter)
   }
 
   /**
    * Called when NetLogo calls clear-all or when the model is reset
    */
   override def clearAll(): Unit = {
     historyLock.synchronized { messageHistory.clear() }
+    usageLock.synchronized {
+      agentUsage.clear()
+      runUsage = UsageTotals.empty
+    }
   }
 
   /**
@@ -182,6 +197,34 @@ class LLMExtension extends DefaultClassManager {
       h += assistant
     }
 
+  /**
+   * Credit a provider reply to the calling agent and to the run.
+   *
+   * Called as soon as a response arrives, before the caller decides whether
+   * the reply is acceptable: tokens were billed either way, so a reply the
+   * extension then rejects (schema parse failure, unmatched choice) still
+   * counts. A call that never produced a response records nothing.
+   */
+  private def recordUsage(agent: Agent, response: ChatResponse): Unit = {
+    val delta = UsageTotals(response.usage.getOrElse(Usage.empty), 1L)
+    usageLock.synchronized {
+      agentUsage.update(agent, agentUsage.getOrElse(agent, UsageTotals.empty).plus(delta))
+      runUsage = runUsage.plus(delta)
+    }
+  }
+
+  private def usageFor(agent: Agent): UsageTotals =
+    usageLock.synchronized { agentUsage.getOrElse(agent, UsageTotals.empty) }
+
+  private def usageForRun: UsageTotals =
+    usageLock.synchronized { runUsage }
+
+  /** The assistant message to store in history, mirroring what chat(messages) returned before. */
+  private def replyMessage(response: ChatResponse): ChatMessage =
+    response.firstMessage.getOrElse(
+      throw new RuntimeException("No response message received from provider")
+    )
+
   /**
    * Get timeout from config, falling back to 30 seconds
    */
@@ -501,8 +544,10 @@ class LLMExtension extends DefaultClassManager {
         val userMessage = ChatMessage.user(inputText)
 
         // Send chat request with user message included, but don't mutate history yet
-        val responseFuture = provider.chat(snapshotHistory(agent) :+ userMessage)
-        val responseMessage = Await.result(responseFuture, getAwaitTimeout)
+        val responseFuture = provider.chatWithFullResponse(snapshotHistory(agent) :+ userMessage)
+        val response = Await.result(responseFuture, getAwaitTimeout)
+        recordUsage(agent, response)
+        val responseMessage = replyMessage(response)
 
         // Only commit both messages after success
         commitExchange(agent, userMessage, responseMessage)
@@ -533,7 +578,9 @@ class LLMExtension extends DefaultClassManager {
 
         // Snapshot on the NetLogo thread; commit the pair atomically on success
         // from the completion thread so overlapping async calls can't interleave.
-        val responseFuture = provider.chat(snapshotHistory(agent) :+ userMessage).map { responseMessage =>
+        val responseFuture = provider.chatWithFullResponse(snapshotHistory(agent) :+ userMessage).map { response =>
+          recordUsage(agent, response)
+          val responseMessage = replyMessage(response)
           commitExchange(agent, userMessage, responseMessage)
           responseMessage.content
         }
@@ -590,8 +637,10 @@ class LLMExtension extends DefaultClassManager {
         tempHistory += userMessage
 
         // Send chat request
-        val responseFuture = provider.chat(tempHistory.toSeq)
-        val responseMessage = Await.result(responseFuture, getAwaitTimeout)
+        val responseFuture = provider.chatWithFullResponse(tempHistory.toSeq)
+        val response = Await.result(responseFuture, getAwaitTimeout)
+        recordUsage(agent, response)
+        val responseMessage = replyMessage(response)
 
         // Commit both template message and response to permanent history on success
         commitExchange(agent, userMessage, responseMessage)
@@ -652,6 +701,7 @@ class LLMExtension extends DefaultClassManager {
         // the constraint behaves exactly as it did before.
         val responseFuture = provider.chatWithFormat(tempHistory.toSeq, EnumFormat(choices))
         val response = Await.result(responseFuture, getAwaitTimeout)
+        recordUsage(agent, response)
 
         // Extract text: prefer content, fall back to thinking field
         val text = response.firstContent.filter(_.nonEmpty)
@@ -751,6 +801,7 @@ class LLMExtension extends DefaultClassManager {
 
         val responseFuture = provider.chatWithFormat(snapshotHistory(agent) :+ userMessage, format)
         val response = Await.result(responseFuture, getAwaitTimeout)
+        recordUsage(agent, response)
 
         val content = response.firstContent.getOrElse("")
 
@@ -811,6 +862,7 @@ class LLMExtension extends DefaultClassManager {
 
         val responseFuture = provider.chatWithFormat(tempHistory.toSeq, JsonObjectFormat)
         val response = Await.result(responseFuture, getAwaitTimeout)
+        recordUsage(agent, response)
 
         val content = response.firstContent.getOrElse("")
 
@@ -893,6 +945,7 @@ class LLMExtension extends DefaultClassManager {
         // Send with user message included, but don't mutate history yet
         val responseFuture = provider.chatWithFullResponse(snapshotHistory(agent) :+ userMessage)
         val response = Await.result(responseFuture, getAwaitTimeout)
+        recordUsage(agent, response)
 
         val answerText = response.firstContent.getOrElse("")
         val thinkingText = response.thinking.getOrElse("")
@@ -1198,6 +1251,44 @@ class LLMExtension extends DefaultClassManager {
     }
   }
 
+  /**
+   * Token accounting as a `[[key value] ...]` list, the same shape structured
+   * output uses so `llm:get` reads it. Counters are numbers; `cost` is a
+   * number only when a provider reported one and `""` otherwise, matching how
+   * JSON null is reported, so an unknown cost is never mistaken for free.
+   */
+  private def usageToLogo(totals: UsageTotals): LogoList = {
+    val u = totals.usage
+    def num(n: Long): AnyRef = Double.box(n.toDouble)
+    LogoList(
+      LogoList("input-tokens", num(u.inputTokens)),
+      LogoList("output-tokens", num(u.outputTokens)),
+      LogoList("total-tokens", num(u.totalTokens)),
+      LogoList("reasoning-tokens", num(u.reasoningTokens.getOrElse(0L))),
+      LogoList("cache-read-tokens", num(u.cacheReadTokens.getOrElse(0L))),
+      LogoList("cache-write-tokens", num(u.cacheWriteTokens.getOrElse(0L))),
+      LogoList("cost", u.cost.map(c => Double.box(c): AnyRef).getOrElse("")),
+      LogoList("latency-ms", num(u.latencyMs.getOrElse(0L))),
+      LogoList("calls", num(totals.calls))
+    )
+  }
+
+  /** llm:usage — token accounting for the calling agent since clear-all. */
+  object UsageReporter extends Reporter {
+    override def getSyntax: Syntax = Syntax.reporterSyntax(ret = Syntax.ListType)
+
+    override def report(args: Array[Argument], context: Context): AnyRef =
+      usageToLogo(usageFor(context.getAgent))
+  }
+
+  /** llm:usage-total — token accounting for every agent in the run since clear-all. */
+  object UsageTotalReporter extends Reporter {
+    override def getSyntax: Syntax = Syntax.reporterSyntax(ret = Syntax.ListType)
+
+    override def report(args: Array[Argument], context: Context): AnyRef =
+      usageToLogo(usageForRun)
+  }
+
   object ActiveReporter extends Reporter {
     override def getSyntax: Syntax = Syntax.reporterSyntax(ret = Syntax.ListType)
 
@@ -1249,3 +1340,13 @@ class LLMExtension extends DefaultClassManager {
     }
   }
 }
+
+/** Running totals: the summed Usage plus how many calls contributed to it. */
+private[llm] case class UsageTotals(usage: Usage, calls: Long) {
+  def plus(other: UsageTotals): UsageTotals =
+    UsageTotals(usage.plus(other.usage), calls + other.calls)
+}
+
+private[llm] object UsageTotals {
+  val empty: UsageTotals = UsageTotals(Usage.empty, 0L)
+}