Capture real token counts from backends and use for compaction (#33)

* Capture real token counts from llama-server and use for compaction

LlamafileClient now records TokenUsage (prompt, completion, total) from
every llama-server response. ContextManager uses these real counts
instead of chars/4 heuristic when available, so compaction and context
warnings fire based on actual KV cache pressure — critical for reasoning
models where think tokens consume context invisibly.

Closes #32

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

* Add token usage capture to OllamaClient

Same pattern as LlamafileClient — captures prompt_eval_count and
eval_count from Ollama responses (both streaming and non-streaming)
into last_usage dict. ContextManager picks these up via the existing
_sync_token_count path in run_inference.

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

---------

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

Author: antoinezambelli

src/forge/__init__.py

@@ -12,7 +12,7 @@
 from forge.core.steps import StepTracker
 from forge.core.inference import InferenceResult, fold_and_serialize, run_inference
 from forge.core.runner import WorkflowRunner
-from forge.clients.base import ChunkType, LLMClient, StreamChunk
+from forge.clients.base import ChunkType, LLMClient, StreamChunk, TokenUsage
 from forge.clients.llamafile import LlamafileClient
 from forge.clients.ollama import OllamaClient
 from forge.context import (
@@ -82,6 +82,7 @@
     "LlamafileClient",
     "OllamaClient",
     "StreamChunk",
+    "TokenUsage",
     # Context
     "CompactEvent",
     "CompactStrategy",

src/forge/clients/base.py

@@ -10,6 +10,21 @@
 from forge.core.workflow import LLMResponse, ToolCall, TextResponse, ToolSpec
 
 
+@dataclass(frozen=True)
+class TokenUsage:
+    """Token counts from a single LLM response.
+
+    Populated from the server's ``usage`` field when available (e.g.
+    llama-server).  Backends that don't report usage leave the client's
+    ``last_usage`` empty and the context manager falls back to heuristic
+    estimation.
+    """
+
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
 # Both Ollama and llama-server use the OpenAI tool schema format today.
 # If a backend diverges, move this back into the relevant client module.
 def format_tool(spec: ToolSpec) -> dict[str, Any]:

src/forge/clients/llamafile.py

@@ -9,7 +9,7 @@
 
 import httpx
 
-from forge.clients.base import ChunkType, StreamChunk, format_tool
+from forge.clients.base import ChunkType, StreamChunk, TokenUsage, format_tool
 from forge.core.workflow import LLMResponse, TextResponse, ToolCall, ToolSpec
 from forge.errors import BackendError, ContextDiscoveryError
 from forge.prompts.templates import build_tool_prompt, extract_tool_call
@@ -144,6 +144,8 @@ def __init__(
         self._cache_prompt = cache_prompt
         self._slot_id = slot_id
 
+        self.last_usage: dict[int, TokenUsage] = {}
+
         if mode in ("native", "prompt"):
             self.resolved_mode: str | None = mode
         else:
@@ -154,6 +156,18 @@ def _apply_slot_id(self, body: dict[str, Any]) -> None:
         if self._slot_id is not None:
             body["slot_id"] = self._slot_id
 
+    def _record_usage(self, data: dict[str, Any]) -> None:
+        """Extract usage from a response and store it keyed by slot ID."""
+        usage = data.get("usage")
+        if not usage:
+            return
+        slot = self._slot_id if self._slot_id is not None else 0
+        self.last_usage[slot] = TokenUsage(
+            prompt_tokens=usage.get("prompt_tokens", 0),
+            completion_tokens=usage.get("completion_tokens", 0),
+            total_tokens=usage.get("total_tokens", 0),
+        )
+
     def _resolve_reasoning(
         self, accumulated_reasoning: str, accumulated_content: str
     ) -> str | None:
@@ -208,6 +222,7 @@ async def send_stream(
             "model": self.model,
             "temperature": self.temperature,
             "stream": True,
+            "stream_options": {"include_usage": True},
             "cache_prompt": self._cache_prompt,
         }
         self._apply_slot_id(body)
@@ -231,6 +246,7 @@ async def send_stream(
 
         accumulated_content = ""
         accumulated_reasoning = ""
+        stream_intentional = False
         # Track multiple tool calls by index — OpenAI streaming sends
         # tool_calls[N] deltas with an index field.
         tool_call_parts: dict[int, dict[str, str]] = {}  # idx -> {name, args}
@@ -252,93 +268,88 @@ async def send_stream(
                 if not line or not line.startswith("data: "):
                     continue
                 data_str = line[6:]
-                # llama-server sends "data: [DONE]" after the final chunk;
-                # llamafile 0.9.x does not — it only sends finish_reason.
                 if data_str == "[DONE]":
-                    stream_done = True
-                else:
-                    chunk = json.loads(data_str)
-                    if "choices" not in chunk or not chunk["choices"]:
-                        continue  # usage-only or error chunk
-                    choice = chunk["choices"][0]
-                    delta = choice.get("delta", {})
-
-                    if "tool_calls" in delta:
-                        for tc_delta in delta["tool_calls"]:
-                            idx = tc_delta.get("index", 0)
-                            if idx not in tool_call_parts:
-                                tool_call_parts[idx] = {"name": "", "args": ""}
-                            func = tc_delta.get("function", {})
-                            if "name" in func:
-                                tool_call_parts[idx]["name"] = func["name"]
-                            if "arguments" in func:
-                                tool_call_parts[idx]["args"] += func["arguments"]
-                                yield StreamChunk(
-                                    type=ChunkType.TOOL_CALL_DELTA,
-                                    content=func["arguments"],
-                                )
-
-                    reasoning_content = delta.get("reasoning_content") or ""
-                    if reasoning_content:
-                        accumulated_reasoning += reasoning_content
-
-                    content = delta.get("content") or ""
-                    if content:
-                        accumulated_content += content
-                        yield StreamChunk(
-                            type=ChunkType.TEXT_DELTA, content=content
-                        )
-
-                    stream_finish_reason = choice.get("finish_reason")
-                    stream_done = stream_finish_reason is not None
-
-                if stream_done:
-                    stream_intentional = stream_finish_reason == "stop"
-                    if tool_call_parts:
-                        reasoning = self._resolve_reasoning(
-                            accumulated_reasoning, accumulated_content
-                        )
-                        result_calls: list[ToolCall] = []
-                        bad_args = False
-                        for idx in sorted(tool_call_parts):
-                            part = tool_call_parts[idx]
-                            try:
-                                args = json.loads(part["args"]) if part["args"] else {}
-                            except json.JSONDecodeError:
-                                bad_args = True
-                                break
-                            result_calls.append(ToolCall(
-                                tool=part["name"],
-                                args=args,
-                                reasoning=reasoning if idx == 0 else None,
-                            ))
-                        if bad_args:
-                            # Model emitted invalid JSON in tool args —
-                            # surface as TextResponse so the runner sends
-                            # a retry nudge instead of crashing.
-                            final: LLMResponse = TextResponse(
-                                content=accumulated_content or part["args"],
-                            )
-                        else:
-                            final = result_calls
-                    elif mode == "prompt" and tools:
-                        think_text, cleaned = _extract_think_tags(
-                            accumulated_content
-                        )
-                        tool_names = [t.name for t in tools]
-                        extracted = extract_tool_call(cleaned, tool_names)
-                        if extracted:
-                            extracted[0].reasoning = self._resolve_reasoning(
-                                accumulated_reasoning, think_text
-                            )
-                            final = extracted
-                        else:
-                            final = TextResponse(content=cleaned, intentional=stream_intentional)
-                    else:
-                        final = TextResponse(content=accumulated_content, intentional=stream_intentional)
-                    yield StreamChunk(type=ChunkType.FINAL, response=final)
                     break
 
+                chunk = json.loads(data_str)
+                if "choices" not in chunk or not chunk["choices"]:
+                    self._record_usage(chunk)
+                    continue
+                choice = chunk["choices"][0]
+                delta = choice.get("delta", {})
+
+                if "tool_calls" in delta:
+                    for tc_delta in delta["tool_calls"]:
+                        idx = tc_delta.get("index", 0)
+                        if idx not in tool_call_parts:
+                            tool_call_parts[idx] = {"name": "", "args": ""}
+                        func = tc_delta.get("function", {})
+                        if "name" in func:
+                            tool_call_parts[idx]["name"] = func["name"]
+                        if "arguments" in func:
+                            tool_call_parts[idx]["args"] += func["arguments"]
+                            yield StreamChunk(
+                                type=ChunkType.TOOL_CALL_DELTA,
+                                content=func["arguments"],
+                            )
+
+                reasoning_content = delta.get("reasoning_content") or ""
+                if reasoning_content:
+                    accumulated_reasoning += reasoning_content
+
+                content = delta.get("content") or ""
+                if content:
+                    accumulated_content += content
+                    yield StreamChunk(
+                        type=ChunkType.TEXT_DELTA, content=content
+                    )
+
+                finish_reason = choice.get("finish_reason")
+                if finish_reason is not None:
+                    stream_intentional = finish_reason == "stop"
+
+            # Stream ended — build and yield FINAL response.
+            if tool_call_parts:
+                reasoning = self._resolve_reasoning(
+                    accumulated_reasoning, accumulated_content
+                )
+                result_calls: list[ToolCall] = []
+                bad_args = False
+                for idx in sorted(tool_call_parts):
+                    part = tool_call_parts[idx]
+                    try:
+                        args = json.loads(part["args"]) if part["args"] else {}
+                    except json.JSONDecodeError:
+                        bad_args = True
+                        break
+                    result_calls.append(ToolCall(
+                        tool=part["name"],
+                        args=args,
+                        reasoning=reasoning if idx == 0 else None,
+                    ))
+                if bad_args:
+                    final: LLMResponse = TextResponse(
+                        content=accumulated_content or part["args"],
+                    )
+                else:
+                    final = result_calls
+            elif mode == "prompt" and tools:
+                think_text, cleaned = _extract_think_tags(
+                    accumulated_content
+                )
+                tool_names = [t.name for t in tools]
+                extracted = extract_tool_call(cleaned, tool_names)
+                if extracted:
+                    extracted[0].reasoning = self._resolve_reasoning(
+                        accumulated_reasoning, think_text
+                    )
+                    final = extracted
+                else:
+                    final = TextResponse(content=cleaned, intentional=stream_intentional)
+            else:
+                final = TextResponse(content=accumulated_content, intentional=stream_intentional)
+            yield StreamChunk(type=ChunkType.FINAL, response=final)
+
     async def get_context_length(self) -> int | None:
         """Query the Llamafile /props endpoint for configured context length.
 
@@ -410,6 +421,7 @@ async def _send_native(
         if resp.status_code != 200:
             raise BackendError(resp.status_code, resp.text)
         data = resp.json()
+        self._record_usage(data)
 
         top_choice = data["choices"][0]
         choice = top_choice["message"]
@@ -470,6 +482,7 @@ async def _send_prompt(
         )
         resp.raise_for_status()
         data = resp.json()
+        self._record_usage(data)
 
         top_choice = data["choices"][0]
         content = top_choice["message"].get("content", "")

src/forge/clients/ollama.py

@@ -8,7 +8,7 @@
 
 import httpx
 
-from forge.clients.base import ChunkType, StreamChunk, format_tool
+from forge.clients.base import ChunkType, StreamChunk, TokenUsage, format_tool
 from forge.core.workflow import LLMResponse, TextResponse, ToolCall, ToolSpec
 from forge.errors import BackendError, ThinkingNotSupportedError
 
@@ -61,6 +61,7 @@ def __init__(
             model_lower = model.lower()
             self._think = any(kw in model_lower for kw in _THINK_HEURISTIC_KEYWORDS)
         self._think_resolved: bool = think is not None
+        self.last_usage: dict[int, TokenUsage] = {}
 
     def _build_options(self) -> dict[str, Any]:
         opts: dict[str, Any] = {"temperature": self.temperature}
@@ -82,6 +83,20 @@ def _resolve_reasoning(
             return None
         return thinking or content or None
 
+    def _record_usage(self, data: dict[str, Any]) -> None:
+        """Extract token usage from an Ollama response."""
+        prompt = data.get("prompt_eval_count")
+        completion = data.get("eval_count")
+        if prompt is None and completion is None:
+            return
+        prompt = prompt or 0
+        completion = completion or 0
+        self.last_usage[0] = TokenUsage(
+            prompt_tokens=prompt,
+            completion_tokens=completion,
+            total_tokens=prompt + completion,
+        )
+
     async def send(
         self,
         messages: list[dict[str, str]],
@@ -118,6 +133,7 @@ async def send(
         if resp.status_code != 200:
             raise BackendError(resp.status_code, resp.text)
         data = resp.json()
+        self._record_usage(data)
 
         if not self._think_resolved:
             self._think_resolved = True
@@ -219,6 +235,7 @@ async def _iter_stream(
                 msg = data.get("message", {})
 
                 if data.get("done"):
+                    self._record_usage(data)
                     tool_calls = msg.get("tool_calls") or pending_tool_calls
                     if tool_calls:
                         reasoning = self._resolve_reasoning(

src/forge/context/manager.py

@@ -88,9 +88,21 @@ def __init__(
         self._context_thresholds = sorted(context_thresholds) if context_thresholds else []
         self._on_context_threshold = on_context_threshold
         self._fired_thresholds: set[float] = set()
+        self._last_known_tokens: int | None = None
+
+    def update_token_count(self, total_tokens: int) -> None:
+        """Record actual token count from the backend.
+
+        Called after each LLM response when the backend reports usage.
+        Subsequent calls to ``estimate_tokens`` will return this value
+        until the next update.
+        """
+        self._last_known_tokens = total_tokens
 
     def estimate_tokens(self, messages: list[Message]) -> int:
-        """Estimate token count via char/4 heuristic."""
+        """Return actual token count if available, else char/4 heuristic."""
+        if self._last_known_tokens is not None:
+            return self._last_known_tokens
         return sum(len(m.content) for m in messages) // 4
 
     def check_thresholds(self, messages: list[Message]) -> str | None:

src/forge/core/inference.py

@@ -13,7 +13,7 @@
 from dataclasses import dataclass, field
 from typing import Any
 
-from forge.clients.base import ChunkType, LLMClient, StreamChunk
+from forge.clients.base import ChunkType, LLMClient, StreamChunk, TokenUsage
 from forge.context.manager import ContextManager
 from forge.core.messages import Message, MessageMeta, MessageRole, MessageType, ToolCallInfo
 from forge.core.workflow import LLMResponse, TextResponse, ToolCall, ToolSpec
@@ -28,6 +28,17 @@
 }
 
 
+def _sync_token_count(client: LLMClient, context_manager: ContextManager) -> None:
+    """Feed actual token count from the client into the context manager."""
+    last_usage = getattr(client, "last_usage", None)
+    if not isinstance(last_usage, dict):
+        return
+    slot_id = getattr(client, "_slot_id", None) or 0
+    usage: TokenUsage | None = last_usage.get(slot_id)
+    if usage is not None:
+        context_manager.update_token_count(usage.total_tokens)
+
+
 @dataclass
 class InferenceResult:
     """Result of a single inference call (may include transparent retries).
@@ -193,6 +204,9 @@ async def run_inference(
         else:
             response = await client.send(api_messages, tools=tool_specs)
 
+        # Update context manager with real token count if available.
+        _sync_token_count(client, context_manager)
+
         # Validate
         validation = validator.validate(response, trust_text_intent=trust_text_intent)