feat: record token counts and the producing model per message

Cost has to survive reopening a conversation, which means storing what
each reply used. The model goes on the message rather than the session
because sessions records only the current one, and a conversation that
switched models would otherwise be priced entirely at whichever is
selected now. All three columns are nullable so old rows read as unknown.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

2 files changed, 127 insertions, 13 deletions

diff --git a/llamachat/db.py b/llamachat/db.py
index 65a2b29..4465adc 100644
--- a/llamachat/db.py
+++ b/llamachat/db.py
@@ -33,13 +33,16 @@ CREATE TABLE IF NOT EXISTS sessions (
 );
 
 CREATE TABLE IF NOT EXISTS messages (
-    id         INTEGER PRIMARY KEY,
-    session_id INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
-    role       TEXT    NOT NULL,
-    content    TEXT    NOT NULL,
-    reasoning  TEXT    NOT NULL DEFAULT '',
-    searches   TEXT,
-    created_at INTEGER NOT NULL
+    id                INTEGER PRIMARY KEY,
+    session_id        INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
+    role              TEXT    NOT NULL,
+    content           TEXT    NOT NULL,
+    reasoning         TEXT    NOT NULL DEFAULT '',
+    searches          TEXT,
+    prompt_tokens     INTEGER,
+    completion_tokens INTEGER,
+    model             TEXT,
+    created_at        INTEGER NOT NULL
 );
 
 CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, id);
@@ -107,6 +110,13 @@ class History:
                 # Nullable rather than defaulted: NULL means "no searches",
                 # which is exactly what every pre-migration row wants.
                 "searches": "TEXT",
+                # Nullable for the same reason: an old row has no counts,
+                # and zero would read as a reply that cost nothing.
+                "prompt_tokens": "INTEGER",
+                "completion_tokens": "INTEGER",
+                # The model on `sessions` is the current one, which prices a
+                # switched conversation wrongly. Record what actually replied.
+                "model": "TEXT",
             },
             "sessions": {
                 "prompt_name": "TEXT NOT NULL DEFAULT ''",
@@ -213,16 +223,23 @@ class History:
         content: str,
         reasoning: str | None = None,
         searches: str | None = None,
+        prompt_tokens: int | None = None,
+        completion_tokens: int | None = None,
+        model: str | None = None,
     ) -> None:
         """Fill in a streamed reply. Omitted fields keep their stored value."""
         columns = ["content = ?"]
         values: list = [content]
-        if reasoning is not None:
-            columns.append("reasoning = ?")
-            values.append(reasoning)
-        if searches is not None:
-            columns.append("searches = ?")
-            values.append(searches)
+        for column, value in (
+            ("reasoning", reasoning),
+            ("searches", searches),
+            ("prompt_tokens", prompt_tokens),
+            ("completion_tokens", completion_tokens),
+            ("model", model),
+        ):
+            if value is not None:
+                columns.append(f"{column} = ?")
+                values.append(value)
         values.append(message_id)
         self.conn.execute(
             f"UPDATE messages SET {', '.join(columns)} WHERE id = ?", values
diff --git a/test_llamachat.py b/test_llamachat.py
index 5f96a80..4d68cc6 100755
--- a/test_llamachat.py
+++ b/test_llamachat.py
@@ -1724,6 +1724,102 @@ def test_metadata_and_cost():
     print("ok  metadata resolution and cost")
 
 
+def test_token_column_migration():
+    """A pre-token database opens, and new rows record counts and model."""
+    import sqlite3
+
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "old.db"
+        conn = sqlite3.connect(path)
+        conn.executescript(
+            "CREATE TABLE sessions (id INTEGER PRIMARY KEY, mode TEXT,"
+            " title TEXT, model TEXT, created_at INTEGER, updated_at INTEGER);"
+            "CREATE TABLE messages (id INTEGER PRIMARY KEY, session_id INTEGER,"
+            " role TEXT NOT NULL, content TEXT NOT NULL,"
+            " created_at INTEGER NOT NULL);"
+            "INSERT INTO sessions VALUES (1,'chat','old','m',0,0);"
+            "INSERT INTO messages VALUES (1,1,'assistant','older reply',0);"
+        )
+        conn.commit()
+        conn.close()
+
+        history = db.History(path)
+        rows = history.messages(1)
+        # The pre-migration row survives and reads as unknown, not as zero.
+        assert rows[0]["content"] == "older reply"
+        assert rows[0]["prompt_tokens"] is None
+        assert rows[0]["completion_tokens"] is None
+        assert rows[0]["model"] is None
+
+        mid = history.add_message(1, "assistant", "")
+        history.update_message(
+            mid, "new reply",
+            prompt_tokens=1200, completion_tokens=340,
+            model="together:Qwen/Qwen2.5",
+        )
+        fresh = history.messages(1)[1]
+        assert fresh["prompt_tokens"] == 1200
+        assert fresh["completion_tokens"] == 340
+        assert fresh["model"] == "together:Qwen/Qwen2.5"
+
+        # Omitting them leaves stored values alone, as with reasoning.
+        history.update_message(mid, "edited")
+        kept = history.messages(1)[1]
+        assert kept["content"] == "edited"
+        assert kept["prompt_tokens"] == 1200
+        assert kept["model"] == "together:Qwen/Qwen2.5"
+        history.close()
+
+    # A fresh database must end up with the same columns as a migrated one,
+    # or cost would read back on one path and raise on the other. This
+    # passes via SCHEMA or via _migrate() indifferently, which is the point:
+    # both paths run on every open and either one alone suffices.
+    from llamachat import models, providers
+
+    with tempfile.TemporaryDirectory() as tmp:
+        fresh_db = db.History(Path(tmp) / "new.db")
+        columns = [
+            row["name"]
+            for row in fresh_db.conn.execute("PRAGMA table_info(messages)")
+        ]
+        assert {"prompt_tokens", "completion_tokens", "model"} <= set(columns)
+        # SCHEMA puts created_at last while ALTER TABLE appends after it, so
+        # the two paths hold the same columns in a different order. That is
+        # tolerable only because every read goes by name: a positional read
+        # of a message row would be right on one path and wrong on the other.
+        assert columns[-1] == "created_at"
+        sid = fresh_db.create_session("chat", "m", "t")
+        mid = fresh_db.add_message(sid, "assistant", "")
+        fresh_db.update_message(
+            mid, "hi", prompt_tokens=7, completion_tokens=3,
+            model="together:Qwen/Qwen2.5",
+        )
+        row = fresh_db.messages(sid)[0]
+        # A stored row prices straight through models.message_cost(), which
+        # proves the three column names are exactly the ones it reads.
+        table = providers.parse(
+            {
+                "providers": {
+                    "together": {
+                        "base_url": "https://api.example.org",
+                        "api_key": "env:X",
+                        "price_in": 1_000_000.0,
+                        "price_out": 1_000_000.0,
+                    }
+                }
+            }
+        )
+        store = models.ModelStore(Path(tmp) / "models.ini")
+        assert models.conversation_cost([row], table, store) == 10.0
+
+        # Reopening runs _migrate() again over columns that already exist.
+        fresh_db.close()
+        again = db.History(Path(tmp) / "new.db")
+        assert again.messages(sid)[0]["prompt_tokens"] == 7
+        again.close()
+    print("ok  token column migration")
+
+
 class _FakeResponse:
     """Enough of an http.client response for urlopen's context manager."""
 
@@ -2469,6 +2565,7 @@ if __name__ == "__main__":
     test_config_providers()
     test_models_store()
     test_metadata_and_cost()
+    test_token_column_migration()
     test_search_tool_schema()
     test_search_results_sanitising()
     test_tool_call_accumulation()