@tiens.nguyen/gonext-local-worker 1.0.249 → 1.0.250

package/gonext-local-worker.mjs

@@ -1980,6 +1980,10 @@ async function runAgentChatJob(job) {
     let inThink = false;
     let finalText = "";
     let answerStreamed = false; // an answer_stream delta already delivered the visible text
+    // Task #83: running total of prompt tokens sent to the agent CODE model. Updated on
+    // each python "tokens" event, persisted to the job doc (Mongo) so the REPL can show
+    // it live and it survives on the completed job.
+    let latestCodeTokens = 0;
     // Raw model-stream tokens (event.type === "stream") are Python code + `Thought:`
     // prose + <code>/<|"|> fragments. The web Thought panel renders reasoning as
     // markdown, which turns `#`-prefixed code comments into <h1> headings and
@@ -2106,6 +2110,19 @@ async function runAgentChatJob(job) {
         }
         enqueueText(event.text);
         answerStreamed = true;
+      } else if (event.type === "tokens" && Number.isFinite(event.codeInput)) {
+        // Prompt tokens sent to the agent CODE model (task #83). Persist the running
+        // total to the job doc so it's in Mongo and the REPL can render it live. Steps
+        // are seconds apart, so a fire-and-forget PATCH per event is cheap; it only
+        // touches agentCodeTokens (never resultText/jobStatus), so it can't race the
+        // chunk/completed writes.
+        if (event.codeInput > latestCodeTokens) {
+          latestCodeTokens = event.codeInput;
+          workerFetch(`/api/worker/jobs/${jobId}`, {
+            method: "PATCH",
+            body: JSON.stringify({ agentCodeTokens: latestCodeTokens }),
+          }).catch(() => {});
+        }
       } else if (event.type === "final" && typeof event.text === "string") {
         closeStreamFence();
         if (inThink) {
@@ -2156,6 +2173,7 @@ async function runAgentChatJob(job) {
         // the whole story.
         resultText: answerStreamed ? fullText : finalText || fullText,
         tokenCount: Math.max(1, Math.ceil((finalText || fullText).length / 4)),
+        ...(latestCodeTokens > 0 ? { agentCodeTokens: latestCodeTokens } : {}),
         totalTimeSeconds,
       }),
     });

package/gonext-repl.mjs

@@ -122,6 +122,12 @@ const SPINNERS = [
 // green, cyan, blue, indigo, pink, amber).
 const WAIT_COLORS_256 = [120, 87, 111, 147, 218, 221];
 const color256 = (n, s) => `\x1b[38;5;${n}m${s}\x1b[0m`;
+// Compact token count: 950 → "950", 12300 → "12.3k", 1500000 → "1.5M" (task #83).
+const fmtTokens = (n) => {
+  if (n < 1000) return String(n);
+  if (n < 1_000_000) return `${(n / 1000).toFixed(n < 100_000 ? 1 : 0)}k`;
+  return `${(n / 1_000_000).toFixed(1)}M`;
+};
 const BULLET = "●"; // a completed action
 
 // ---------- flags ----------
@@ -1065,6 +1071,9 @@ async function runAgentTurn(history) {
   // poll figure out how much of the FINAL answer is still unshown even when resultText
   // was replaced (not extended) on completion. See the completed-branch recovery below.
   let liveAnswer = "";
+  // Running total of prompt tokens sent to the agent CODE model this turn (task #83),
+  // read from the job doc each poll. Rendered on the thinking line after the playful word.
+  let latestCodeTokens = 0;
   // First few KB of the raw stream as consumed — a cheap fingerprint to verify the
   // completed resultText really EXTENDS what we streamed (an old worker replaces it
   // with the shorter bare answer, which must not be sliced by shownChars).
@@ -1139,9 +1148,12 @@ async function runAgentTurn(history) {
     // Line 2 (indented): the SPINNER + the playful word together, both in the cycling
     // color — the spinner belongs with the word (dot on the label line, spinner leading
     // the flavor line).
+    // Task #83: after the playful word, show the running agent-code-model token count
+    // (dim), e.g. "  ⠙ Caffeinating… · 12.3k tok". Only once we've seen some tokens.
+    const tokLabel = latestCodeTokens > 0 ? dim(` · ${fmtTokens(latestCodeTokens)} tok`) : "";
     process.stdout.write(
       `${green(BULLET)} ${green(`${fit(primary)}… (${secs}s)`)}` +
-        "\n" + `  ${color256(wc, `${glyph} ${thinkWord}…`)}`
+        "\n" + `  ${color256(wc, `${glyph} ${thinkWord}…`)}${tokLabel}`
     );
     statusShown = true;
   };
@@ -1407,6 +1419,10 @@ async function runAgentTurn(history) {
         continue;
       }
       jobRunning = job.jobStatus === "running";
+      // Live agent-code-model token total (task #83) — monotonic, shown on the thinking line.
+      if (Number.isFinite(job.agentCodeTokens) && job.agentCodeTokens > latestCodeTokens) {
+        latestCodeTokens = job.agentCodeTokens;
+      }
       const text = typeof job.resultText === "string" ? job.resultText : "";
       if (job.jobStatus === "completed") {
         // The job can flip to "completed" between polls with authoritative resultText that

package/gonext_agent_chat.py

@@ -3936,6 +3936,10 @@ def run_agent_chat(cfg):
     # literal messages array sent to /v1/chat/completions.
     class _LoggingModel(OpenAIServerModel):
         def __init__(self, *a, stream_agent=False, **kw):
+            # Running total of PROMPT (input) tokens sent to the CODE model across the
+            # turn (task #83). Only this wrapper counts — the chat/brain helper calls
+            # (_chat_completion) don't route through here, so this is "code model only".
+            self._code_tokens_total = 0
             # stream_agent: request the completion with stream=True and reassemble the
             # deltas into one message. Enabled for Ollama, where a long single-shot
             # (non-stream) generation sends nothing until the very end — a reverse proxy
@@ -4129,6 +4133,29 @@ def run_agent_chat(cfg):
                 token_usage=TokenUsage(input_tokens=in_tok, output_tokens=out_tok),
             )
 
+        def _estimate_prompt_tokens(self, args, kwargs):
+            """Fallback token estimate (~4 chars/token) from the outgoing messages, used
+            only when the model server doesn't report prompt usage (task #83)."""
+            msgs = kwargs.get("messages")
+            if msgs is None and args:
+                msgs = args[0]
+            if not msgs:
+                return 0
+            chars = 0
+            try:
+                for m in msgs:
+                    c = m.get("content") if isinstance(m, dict) else getattr(m, "content", None)
+                    if isinstance(c, list):  # multimodal content parts
+                        for part in c:
+                            t = part.get("text") if isinstance(part, dict) else None
+                            if isinstance(t, str):
+                                chars += len(t)
+                    elif isinstance(c, str):
+                        chars += len(c)
+            except Exception:  # noqa: BLE001
+                return 0
+            return max(1, chars // 4)
+
         def generate(self, *args, **kwargs):
             # Two safeguards on each step's reply BEFORE smolagents parses it for the
             # code block / final_answer:
@@ -4199,6 +4226,23 @@ def run_agent_chat(cfg):
                     raise last_err
             finally:
                 hb_stop.set()
+            # Task #83: tally PROMPT tokens sent to the code model this step, then emit
+            # the running total so the worker can persist it (Mongo) and the REPL can
+            # show it live. Best-effort — never let accounting break the run.
+            try:
+                tu = getattr(msg, "token_usage", None)
+                step_in = int(getattr(tu, "input_tokens", 0) or 0) if tu else 0
+                if step_in <= 0:
+                    step_in = self._estimate_prompt_tokens(args, kwargs)
+                if step_in > 0:
+                    self._code_tokens_total += step_in
+                    _emit({
+                        "type": "tokens",
+                        "codeInput": self._code_tokens_total,
+                        "stepInput": step_in,
+                    })
+            except Exception:  # noqa: BLE001
+                pass
             try:
                 content = getattr(msg, "content", None)
                 if content is None:

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tiens.nguyen/gonext-local-worker",
-  "version": "1.0.249",
+  "version": "1.0.250",
   "description": "Polls GoNext cloud API for async local LLM jobs and runs them against Ollama/OpenAI-compatible servers on this Mac",
   "type": "module",
   "license": "MIT",