1 files changed, 97 insertions, 0 deletions
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()