3 files changed, 79 insertions, 8 deletions

diff --git a/tests/test_clear_messages.py b/tests/test_clear_messages.py
index 0b25cff..c3f3595 100644
--- a/tests/test_clear_messages.py
+++ b/tests/test_clear_messages.py
@@ -58,4 +58,36 @@ def test_clear_messages_invalid_session_id():
 
 def test_clear_messages_missing_session():
     store = make_store()
-    assert store.clear_messages("0" * 32) is None
\ No newline at end of file
+    assert store.clear_messages("0" * 32) is None
+
+
+def test_add_message_stores_token_count_and_clears_it():
+    store = make_store()
+    session = store.create(title="test")
+    session_id = session["id"]
+    store.add_message(
+        session_id=session_id,
+        role="assistant",
+        content="hello",
+        token_count={
+            "prompt_eval_count": 120,
+            "eval_count": 80,
+            "num_ctx": 8192,
+            "total_count": 200,
+        },
+    )
+    reloaded = store.get(session_id)
+    assert reloaded is not None
+    assert reloaded["messages"][0]["token_count"]["prompt_eval_count"] == 120
+    assert reloaded["token_count"] == {
+        "prompt_eval_count": 120,
+        "eval_count": 80,
+        "total_count": 200,
+    }
+
+    cleared = store.clear_messages(session_id)
+    assert cleared is not None
+    assert "token_count" not in cleared
+    reloaded_after = store.get(session_id)
+    assert reloaded_after is not None
+    assert "token_count" not in reloaded_after
\ No newline at end of file
diff --git a/tests/test_llm.py b/tests/test_llm.py
index cbb9910..3002a77 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -12,6 +12,7 @@ from ollama._types import ResponseError
 from ollama_web.llm import (
     _embed_images_in_content,
     _is_image_message,
+    _token_counts,
     chat_with_tools,
     stream_chat_with_tools,
 )
@@ -81,3 +82,26 @@ def test_stream_chat_with_tools_accepts_capability_set() -> None:
     )
     with pytest.raises(ResponseError):
         list(gen)
+
+
+def test_token_counts_extracts_from_dict() -> None:
+    assert _token_counts({"prompt_eval_count": 100, "eval_count": 50}) == {
+        "prompt_eval_count": 100,
+        "eval_count": 50,
+    }
+
+
+def test_token_counts_extracts_from_object() -> None:
+    class FakeResp:
+        prompt_eval_count = 200
+        eval_count = 30
+
+    assert _token_counts(FakeResp()) == {
+        "prompt_eval_count": 200,
+        "eval_count": 30,
+    }
+
+
+def test_token_counts_returns_none_when_missing() -> None:
+    assert _token_counts({}) is None
+    assert _token_counts(None) is None
\ No newline at end of file
diff --git a/tests/test_models.py b/tests/test_models.py
index 8cc2738..177b86b 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from unittest.mock import patch
+
 from starlette.testclient import TestClient
 
 from ollama_web.app import create_app
@@ -14,6 +16,17 @@ def _login_client(client: TestClient) -> str:
     return str(login.headers["x-csrf-token"])
 
 
+class _FakeClient:
+    def __init__(self, capabilities: list[str] | None = None) -> None:
+        self._capabilities = capabilities or []
+
+    def show(self, _model: str) -> dict[str, list[str]]:
+        return {"capabilities": self._capabilities}
+
+    def list(self) -> dict[str, list[dict[str, str]]]:
+        return {"models": []}
+
+
 def test_model_capabilities_route_supports_slash_in_model_name() -> None:
     """HuggingFace-style model names contain slashes and must reach the API.
 
@@ -21,11 +34,12 @@ def test_model_capabilities_route_supports_slash_in_model_name() -> None:
     so the route uses the ``path`` converter. Encoding ``%2F`` must decode to
     the full model name, not split the URL into multiple segments.
     """
-    with TestClient(create_app()) as client:
-        csrf = _login_client(client)
+    slash_model = "hf.co/hotdogs/gemma-4-E4B-it-ultra-uncensored-heretic-GGUF:Q5_K_M"
+    encoded = slash_model.replace("/", "%2F").replace(":", "%3A")
+    fake = _FakeClient(["completion", "vision", "tools", "thinking"])
 
-        slash_model = "hf.co/hotdogs/gemma-4-E4B-it-ultra-uncensored-heretic-GGUF:Q5_K_M"
-        encoded = slash_model.replace("/", "%2F").replace(":", "%3A")
+    with patch("ollama_web.llm.get_client", return_value=fake), TestClient(create_app()) as client:
+        csrf = _login_client(client)
         response = client.get(
             f"/api/models/{encoded}/capabilities",
             headers={"X-CSRF-Token": csrf},
@@ -41,9 +55,10 @@ def test_model_capabilities_route_supports_slash_in_model_name() -> None:
 
 def test_model_capabilities_route_still_works_without_slashes() -> None:
     """Model names without slashes (e.g. cloud aliases) must keep working."""
-    with TestClient(create_app()) as client:
-        csrf = _login_client(client)
+    fake = _FakeClient(["completion"])
 
+    with patch("ollama_web.llm.get_client", return_value=fake), TestClient(create_app()) as client:
+        csrf = _login_client(client)
         response = client.get(
             "/api/models/glm-5.2%3Acloud/capabilities",
             headers={"X-CSRF-Token": csrf},
@@ -51,4 +66,4 @@ def test_model_capabilities_route_still_works_without_slashes() -> None:
         assert response.status_code == 200
         payload = response.json()
         assert "capabilities" in payload
-        assert "completion" in [c.lower() for c in payload["capabilities"]]
+        assert "completion" in [c.lower() for c in payload["capabilities"]]
\ No newline at end of file