add token usage

AGENTS.md

@@ -11,12 +11,12 @@ This file is the top-level orientation map. Depth lives under [docs/](docs/).
 
 | Path | Purpose |
 | --- | --- |
-| [oli_bot/chat.py](oli_bot/chat.py) | Textual TUI app (`OliBot`). Owns command handling, session UI, the `#command-suggestions` autocomplete `ListView`, and (when pooling is enabled) the "Active Sub-Agents" `Tree`. Slash-command names come from the module-level `COMMANDS` tuple. |
+| [oli_bot/chat.py](oli_bot/chat.py) | Textual TUI app (`OliBot`). Owns command handling, session UI, the `#command-suggestions` autocomplete `ListView`, and (when pooling is enabled) the "Active Sub-Agents" `Tree`. Tracks cumulative session token usage (persisted per session) and renders it in the `#status-bar`. Slash-command names come from the module-level `COMMANDS` tuple. |
 | [oli_bot/api_server.py](oli_bot/api_server.py) | FastAPI app exposing the harness over an OpenAI-compatible REST API (`GET /v1/models`, `POST /v1/chat/completions` streaming + non-streaming, `GET /health`). Stateless from the caller's POV; a single process-private `Agent` is shared across requests and serialised with a `threading.RLock`. Auto-approves permissions (no human), but offline/dry-run still apply. |
-| [oli_bot/agent.py](oli_bot/agent.py) | `Agent` — mode + system prompt owner; orchestrates the tool-calling loop and streams typed events (`TextChunk`, `ThinkingChunk`, `ToolCallChunk`, `ToolCallExecuting`, `ToolCallResult`, `StreamChunk`, `Error`, `Done`). Also hosts `sanitize_tool_history`, `stream_sub_agent_run`, and `AgentPool` (built from [oli_bot/agents.yaml](oli_bot/agents.yaml) when `--use-pool` is set). |
-| [oli_bot/backends/](oli_bot/backends/) | Backend package — `ModelBackend` ABC, `OllamaBackend`, `OpenAIBackend`, `HuggingFaceBackend`, `TransformersBackend`, and the `create_model_backend()` factory. Also hosts the shared `_StreamingThinkParser` and per-backend message formatting (Ollama native `images`, OpenAI `image_url` or Bedrock-native blocks via `openai_vision_style`, textual placeholder for text-only backends). See [docs/BACKENDS.md](docs/BACKENDS.md). |
+| [oli_bot/agent.py](oli_bot/agent.py) | `Agent` — mode + system prompt owner; orchestrates the tool-calling loop and streams typed events (`TextChunk`, `ThinkingChunk`, `ToolCallChunk`, `ToolCallExecuting`, `ToolCallResult`, `StreamChunk`, `UsageEvent`, `Error`, `Done`). Aggregates per-call `UsageChunk`s from each backend round into a single per-run `UsageEvent`. Also hosts `sanitize_tool_history`, `_merge_usage`, `stream_sub_agent_run`, and `AgentPool` (built from [oli_bot/agents.yaml](oli_bot/agents.yaml) when `--use-pool` is set). |
+| [oli_bot/backends/](oli_bot/backends/) | Backend package — `ModelBackend` ABC, `OllamaBackend`, `OpenAIBackend`, `HuggingFaceBackend`, `TransformersBackend`, and the `create_model_backend()` factory. Also hosts the shared `_StreamingThinkParser` and per-backend message formatting (Ollama native `images`, OpenAI `image_url` or Bedrock-native blocks via `openai_vision_style`, textual placeholder for text-only backends). Every backend surfaces a trailing `UsageChunk`: exact counts from provider usage where available (OpenAI `usage`/`stream_options`, Ollama `prompt_eval_count`/`eval_count`, HF `usage`), else a `~chars/4` estimate via `estimate_tokens`. See [docs/BACKENDS.md](docs/BACKENDS.md). |
 | [oli_bot/screens/](oli_bot/screens/) | All `ModalScreen` subclasses: `PermissionScreen`, `ConfirmScreen`, `ModelPickerScreen`, `ServerListScreen`, `MCPSetupScreen`, `SessionListScreen`, `WorkspaceListScreen`, `SubAgentViewScreen`, `ConfigScreen`, `InputPromptScreen`, plus `taglines.py` / `todo_widget.py`. |
-| [oli_bot/models.py](oli_bot/models.py) | Shared dataclasses: `Message`, `ToolCall`, `ModelResponse`, `HostConfig`, `MCPServerConfig`, `ProfileData`, `SubAgentRun`, `ImageAttachment`, `TodoItem` / `TodoListState`, plus `AgentEvent` variants and the `AgentRole` enum. `Message.images` is in-memory only (dropped on session save). |
+| [oli_bot/models.py](oli_bot/models.py) | Shared dataclasses: `Message`, `ToolCall`, `ModelResponse`, `HostConfig`, `MCPServerConfig`, `ProfileData`, `SubAgentRun`, `ImageAttachment`, `TodoItem` / `TodoListState`, plus `AgentEvent` variants and the `AgentRole` enum. Token accounting lives here too: `Usage` (prompt/completion/`estimated` flag), the per-call `UsageChunk` stream event, and the per-run `UsageEvent`. `Message.images` is in-memory only (dropped on session save); `ModelResponse.usage` is optionally set by backends. |
 | [oli_bot/config.py](oli_bot/config.py) | `AppConfig` — `pydantic_settings.BaseSettings`. Env vars prefixed `OLI_`, plus `.env` support and `OLI_TRUNCATION_SMALL` / `_LARGE` aliases via `AliasChoices`. Module-level `configs = AppConfig()` singleton. See [docs/CONFIGURE.md](docs/CONFIGURE.md). |
 | [oli_bot/settings.py](oli_bot/settings.py) | `SettingsManager` — load/save/merge `~/.config/oli/settings.json`; precedence `settings.json` > `OLI_*` env > SDK-standard env (`OPENAI_API_KEY`, `OPENAI_BASE_URL`, `HUGGINGFACE_API_KEY`, `HF_TOKEN`) > declared defaults. Empty API-key strings in JSON fall through to env. |
 | [oli_bot/profiles/](oli_bot/profiles/) | `ProfileManifest` / `PermissionsManifest` (Pydantic) in `schema.py`; `ProfilePermissionEnforcer` (layered allow/deny glob patterns, base-profile inheritance, deny-overrides-allow) in `permissions.py`. Built-in profiles ship as sibling directories. |
@@ -69,6 +69,7 @@ See [docs/AGENT-POOLING.md](docs/AGENT-POOLING.md) for the `agents.yaml` schema,
 
 - **Session lifecycle** — new UUID by default; `-s`/`--load-session <uuid>` loads a specific one, `--resume-last` loads the newest for the active server (mutually exclusive). Auto-saved after every response; `/sessions` opens `SessionListScreen` for browse/switch/rename/delete. On exit `main()` prints the resume hint.
 - **Tool loop** — `Agent._tool_loop` calls `backend.stream_generate()`, streaming text and thinking to the TUI while accumulating tool calls; after each round it calls `MCPClientManager.drain_builtin_attachments()` so `view_image` results ride into the next turn as a synthetic user-role `Message(images=...)` (kept **after** all `role=tool` messages so the assistant→tool block stays contiguous for `sanitize_tool_history`). If iterations are exhausted, `_final_stream` does one text-only pass to guarantee a reply.
+- **Token accounting** — each backend yields a trailing `UsageChunk` (exact where the provider reports counts, else `~chars/4` via `estimate_tokens`). The agent sums per-call chunks across a run and emits a single `UsageEvent` before `Done`; the TUI adds it to the session total in the `#status-bar` (`~` prefix when any value was estimated). The running total persists in session JSON (`total_tokens` / `total_tokens_estimated`) and is reset on `/clear`, `/home`, and new-session flows.
 - **Error paths** — a failed round yields `Error` for the UI then `Done(full_text="")`; the empty `full_text` deliberately trips the TUI's "skip empty assistant append" branch so error text never enters `self.messages` (prevents the Bedrock/Anthropic 400 poisoning loop). All `stream_generate()` implementations re-raise instead of swallowing.
 - **Permission gating** — `BuiltinToolManager.call_tool()` runs the profile enforcer, then `Session.needs_permission()`, then (if needed) `confirm_callback(description)`; the TUI shows `PermissionScreen` and the callback returns `"once"`, `"session"`, or `"deny"`. Session grants persist per scope for the process lifetime. `glob`/`grep` targeting patterns like `.env*`, `*.pem`, `*secret*` trigger the `workspace_sensitive` scope even inside the workspace.
 - **Built-in tool naming** — registered as `builtin__<name>` and dispatched by `MCPClientManager.call_tool`.
@@ -147,7 +148,7 @@ pip install -e '.[dev]'
 pytest
 ```
 
-Tests under `tests/` cover: `AgentPool` scaffolding (lookup, `${VAR}` expansion, `root-agent` exclusion, per-agent backend overrides, concurrent dispatch), config env-var precedence, session round-trip + `save_session` ID propagation, permission matrix (workspace scoping, sensitive files, `glob`/`grep` sensitive-pattern gating), truncation boundary preservation, security regressions (`git` restricted to read-only subcommands in `run_command`, `find -exec/-delete` blocked, SSRF for loopback/private/link-local/non-`http(s)`), `OpenAIBackend.stream_generate` tool-call flushing on all finish reasons, and the `api_server.py` OpenAI-compatible endpoints. `tests/conftest.py` scrubs `OLI_*` env vars so runs are hermetic.
+Tests under `tests/` cover: `AgentPool` scaffolding (lookup, `${VAR}` expansion, `root-agent` exclusion, per-agent backend overrides, concurrent dispatch), config env-var precedence, session round-trip + `save_session` ID propagation, permission matrix (workspace scoping, sensitive files, `glob`/`grep` sensitive-pattern gating), truncation boundary preservation, security regressions (`git` restricted to read-only subcommands in `run_command`, `find -exec/-delete` blocked, SSRF for loopback/private/link-local/non-`http(s)`), `OpenAIBackend.stream_generate` tool-call flushing on all finish reasons, token-usage accounting (backend `UsageChunk` exact/estimate paths, agent `UsageEvent` aggregation, session `total_tokens` round-trip), and the `api_server.py` OpenAI-compatible endpoints. `tests/conftest.py` scrubs `OLI_*` env vars so runs are hermetic.
 
 ## Code style
 

README.md

@@ -44,6 +44,7 @@ Concretely, that means:
 - **Agent pooling (optional)** — with `--use-pool`, the root agent can fan tasks out concurrently to vendor-agnostic sub-agents defined in [agents.yaml](agents.yaml) via a `dispatch` tool. Each pool entry binds a model _and_ a backend, so dispatch decisions are also compute-location decisions — a frontier model can plan while sensitive work stays on a local model, or a local root can fan out to faster remote SLMs for latency-sensitive tool calls.
 - **Permission system** — write operations and sensitive reads require user approval. Session grants, workspace scoping, and profile-level allow/deny lists.
 - **Streaming Markdown** responses in a Textual TUI, with in-app model/server/session/profile management.
+- **Session token counter** — the TUI status bar tracks cumulative session token usage (exact where the provider reports counts, `~`-prefixed estimates otherwise), persisted per session and summed from OpenAI/`stream_options`, Ollama eval counts, HuggingFace usage, and Transformers runs. Reset on `/clear`, `/home`, and new-session flows.
 - **OpenAI-compatible API server** — run the same agent harness behind `/v1/models` and `/v1/chat/completions` (streaming + non-streaming) so any workflow that speaks the OpenAI wire protocol (the `openai` Python SDK, curl, or plain REST) can drive the agent.
 
 ## Roadmap / areas of active exploration

oli_bot/agent.py

@@ -30,6 +30,9 @@
     StreamChunk,
     SubAgentRun,
     ThinkingChunk,
+    Usage,
+    UsageChunk,
+    UsageEvent,
 )
 
 logger = logging.getLogger(__name__)
@@ -41,6 +44,7 @@
     | AssistantResponse
     | StreamChunk
     | ThinkingChunk
+    | UsageEvent
     | Error
     | Done
 )
@@ -113,6 +117,15 @@ def sanitize_tool_history(messages: List[Message]) -> List[Message]:
     return kept
 
 
+def _merge_usage(acc: Usage, u: Usage) -> Usage:
+    """Sum per-call usage into a cumulative run total (flag estimated if any)."""
+    return Usage(
+        prompt_tokens=acc.prompt_tokens + u.prompt_tokens,
+        completion_tokens=acc.completion_tokens + u.completion_tokens,
+        estimated=acc.estimated or u.estimated,
+    )
+
+
 class Agent:
     def __init__(
         self,
@@ -287,6 +300,7 @@ async def _tool_loop(
         tools: Optional[List[Dict[str, Any]]],
         confirm_callback: Optional[Callable[[str], Any]],
     ) -> AsyncIterator[AgentEvent]:
+        run_usage = Usage()
         for _ in range(self.config.max_tool_iterations):
             stream = self.backend.stream_generate(
                 messages, tools=tools if tools else []
@@ -310,6 +324,8 @@ async def _tool_loop(
                         yield StreamChunk(event.text)
                     elif isinstance(event, ThinkingChunk):
                         yield event
+                    elif isinstance(event, UsageChunk):
+                        run_usage = _merge_usage(run_usage, event.usage)
                     elif isinstance(event, ToolCallChunk):
                         tool_calls = event.tool_calls
                         for tc in tool_calls:
@@ -344,6 +360,8 @@ async def _tool_loop(
                 return
 
             if not tool_calls:
+                if run_usage.total_tokens:
+                    yield UsageEvent(usage=run_usage)
                 yield Done(full_text=full_response)
                 return
 
@@ -410,6 +428,7 @@ async def _final_stream(
     ) -> AsyncIterator[AgentEvent]:
         full_response = ""
         errored = False
+        usage = Usage()
         try:
             # Forward tools so backends that require a matching tool schema
             # for any toolUse/toolResult blocks in history (e.g. Bedrock via
@@ -428,6 +447,8 @@ async def _final_stream(
                     yield StreamChunk(event.text)
                 elif isinstance(event, ThinkingChunk):
                     yield event
+                elif isinstance(event, UsageChunk):
+                    usage = _merge_usage(usage, event.usage)
         except asyncio.TimeoutError:
             err_text = (
                 f"Response timed out (no data for "
@@ -448,6 +469,8 @@ async def _final_stream(
             else:
                 if not full_response:
                     full_response = "The model completed all tool calls but did not produce a final summary."
+                if usage.total_tokens:
+                    yield UsageEvent(usage=usage)
                 yield Done(full_text=full_response)
 
 

oli_bot/backends/base.py

@@ -11,6 +11,7 @@
     TextChunk,
     ThinkingChunk,
     ToolCallChunk,
+    UsageChunk,
 )
 
 logger = logging.getLogger(__name__)
@@ -20,7 +21,16 @@
 MAX_TOKENS = configs.max_tokens
 TEMPERATURE = configs.temperature
 
-StreamEvent = TextChunk | ToolCallChunk | ThinkingChunk
+StreamEvent = TextChunk | ToolCallChunk | ThinkingChunk | UsageChunk
+
+
+def estimate_tokens(text: str) -> int:
+    """Rough token estimate for providers that do not report usage.
+
+    Mirrors the heuristic used by the OpenAI-compatible API layer (~4 chars per
+    token), so numbers stay consistent everywhere estimation is needed.
+    """
+    return max(1, (len(text) + 3) // 4)
 
 
 class ModelBackend(ABC):
@@ -51,6 +61,7 @@ async def stream_generate(
     "MAX_TOKENS",
     "TEMPERATURE",
     "StreamEvent",
+    "estimate_tokens",
     "logger",
     "Any",
     "AsyncIterator",

oli_bot/backends/huggingface.py

@@ -13,9 +13,17 @@
     ThinkingChunk,
     ToolCall,
     ToolCallChunk,
+    Usage,
+    UsageChunk,
 )
 
-from .base import MAX_TOKENS, TEMPERATURE, ModelBackend, StreamEvent
+from .base import (
+    MAX_TOKENS,
+    TEMPERATURE,
+    ModelBackend,
+    StreamEvent,
+    estimate_tokens,
+)
 from .messages import _format_messages, _format_tools
 from .streaming import _StreamingThinkParser
 
@@ -86,10 +94,29 @@ async def generate(
                     )
                     for i, tc in enumerate(message.tool_calls)
                 ]
+            usage = None
+            usage_field = getattr(response, "usage", None)
+            if (
+                usage_field is not None
+                and getattr(usage_field, "prompt_tokens", None) is not None
+            ):
+                usage = Usage(
+                    prompt_tokens=usage_field.prompt_tokens,
+                    completion_tokens=getattr(usage_field, "completion_tokens", 0) or 0,
+                )
+            else:
+                usage = Usage(
+                    prompt_tokens=estimate_tokens(
+                        json.dumps(_format_messages(messages, stringify_arguments=True))
+                    ),
+                    completion_tokens=estimate_tokens(message.content or ""),
+                    estimated=True,
+                )
             return ModelResponse(
                 content=message.content or "",
                 tool_calls=tool_calls or None,
                 finish_reason=response.choices[0].finish_reason or "stop",
+                usage=usage,
             )
         except Exception as e:
             logger.exception(
@@ -113,13 +140,17 @@ async def stream_generate(
             tool_calls_acc: Dict[int, dict] = {}
             flushed = False
             parser = _StreamingThinkParser()
+            streamed_text = ""
+            chunk_usage = None
             async for chunk in stream:
+                chunk_usage = getattr(chunk, "usage", None) or chunk_usage
                 if chunk.choices and chunk.choices[0].delta:
                     delta = chunk.choices[0].delta
                     reasoning = getattr(delta, "reasoning_content", None)
                     if reasoning:
                         yield ThinkingChunk(reasoning)
                     if delta.content:
+                        streamed_text += delta.content
                         for kind, text in parser.feed(delta.content):
                             if text:
                                 yield (
@@ -192,6 +223,21 @@ async def stream_generate(
                     for data in tool_calls_acc.values()
                 ]
                 yield ToolCallChunk(tool_calls)
+
+            if chunk_usage is not None:
+                usage = Usage(
+                    prompt_tokens=getattr(chunk_usage, "prompt_tokens", 0) or 0,
+                    completion_tokens=getattr(chunk_usage, "completion_tokens", 0) or 0,
+                )
+            else:
+                usage = Usage(
+                    prompt_tokens=estimate_tokens(
+                        json.dumps(_format_messages(messages, stringify_arguments=True))
+                    ),
+                    completion_tokens=estimate_tokens(streamed_text),
+                    estimated=True,
+                )
+            yield UsageChunk(usage)
         except Exception as e:
             logger.exception(
                 "Error in %s.stream_generate: %s", self.__class__.__name__, str(e)