fix(insights): persist token usage for non-CLI sessions
This commit is contained in:
@@ -825,43 +825,6 @@ class TestLastPromptTokens:
|
||||
store.update_session("k1", last_prompt_tokens=0)
|
||||
assert entry.last_prompt_tokens == 0
|
||||
|
||||
def test_update_session_passes_model_to_db(self, tmp_path):
|
||||
"""Gateway session updates should forward the resolved model to SQLite."""
|
||||
config = GatewayConfig()
|
||||
with patch("gateway.session.SessionStore._ensure_loaded"):
|
||||
store = SessionStore(sessions_dir=tmp_path, config=config)
|
||||
store._loaded = True
|
||||
store._save = MagicMock()
|
||||
store._db = MagicMock()
|
||||
|
||||
from gateway.session import SessionEntry
|
||||
from datetime import datetime
|
||||
entry = SessionEntry(
|
||||
session_key="k1",
|
||||
session_id="s1",
|
||||
created_at=datetime.now(),
|
||||
updated_at=datetime.now(),
|
||||
)
|
||||
store._entries = {"k1": entry}
|
||||
|
||||
store.update_session("k1", model="openai/gpt-5.4")
|
||||
|
||||
store._db.set_token_counts.assert_called_once_with(
|
||||
"s1",
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
cache_read_tokens=0,
|
||||
cache_write_tokens=0,
|
||||
estimated_cost_usd=None,
|
||||
cost_status=None,
|
||||
cost_source=None,
|
||||
billing_provider=None,
|
||||
billing_base_url=None,
|
||||
model="openai/gpt-5.4",
|
||||
absolute=True,
|
||||
)
|
||||
|
||||
|
||||
class TestRewriteTranscriptPreservesReasoning:
|
||||
"""rewrite_transcript must not drop reasoning fields from SQLite."""
|
||||
|
||||
|
||||
@@ -126,15 +126,5 @@ async def test_handle_message_persists_agent_token_counts(monkeypatch):
|
||||
assert result == "ok"
|
||||
runner.session_store.update_session.assert_called_once_with(
|
||||
session_entry.session_key,
|
||||
input_tokens=120,
|
||||
output_tokens=45,
|
||||
cache_read_tokens=0,
|
||||
cache_write_tokens=0,
|
||||
last_prompt_tokens=80,
|
||||
model="openai/test-model",
|
||||
estimated_cost_usd=None,
|
||||
cost_status=None,
|
||||
cost_source=None,
|
||||
provider=None,
|
||||
base_url=None,
|
||||
)
|
||||
|
||||
62
tests/test_token_persistence_non_cli.py
Normal file
62
tests/test_token_persistence_non_cli.py
Normal file
@@ -0,0 +1,62 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from run_agent import AIAgent
|
||||
|
||||
|
||||
def _mock_response(*, usage: dict, content: str = "done"):
|
||||
msg = SimpleNamespace(content=content, tool_calls=None)
|
||||
choice = SimpleNamespace(message=msg, finish_reason="stop")
|
||||
return SimpleNamespace(
|
||||
choices=[choice],
|
||||
model="test/model",
|
||||
usage=SimpleNamespace(**usage),
|
||||
)
|
||||
|
||||
|
||||
def _make_agent(session_db, *, platform: str):
|
||||
with (
|
||||
patch("run_agent.get_tool_definitions", return_value=[]),
|
||||
patch("run_agent.check_toolset_requirements", return_value={}),
|
||||
patch("run_agent.OpenAI"),
|
||||
):
|
||||
agent = AIAgent(
|
||||
api_key="test-key",
|
||||
quiet_mode=True,
|
||||
skip_context_files=True,
|
||||
skip_memory=True,
|
||||
session_db=session_db,
|
||||
session_id=f"{platform}-session",
|
||||
platform=platform,
|
||||
)
|
||||
agent.client = MagicMock()
|
||||
agent.client.chat.completions.create.return_value = _mock_response(
|
||||
usage={
|
||||
"prompt_tokens": 11,
|
||||
"completion_tokens": 7,
|
||||
"total_tokens": 18,
|
||||
}
|
||||
)
|
||||
return agent
|
||||
|
||||
|
||||
def test_run_conversation_persists_tokens_for_telegram_sessions():
|
||||
session_db = MagicMock()
|
||||
agent = _make_agent(session_db, platform="telegram")
|
||||
|
||||
result = agent.run_conversation("hello")
|
||||
|
||||
assert result["final_response"] == "done"
|
||||
session_db.update_token_counts.assert_called_once()
|
||||
assert session_db.update_token_counts.call_args.args[0] == "telegram-session"
|
||||
|
||||
|
||||
def test_run_conversation_persists_tokens_for_cron_sessions():
|
||||
session_db = MagicMock()
|
||||
agent = _make_agent(session_db, platform="cron")
|
||||
|
||||
result = agent.run_conversation("hello")
|
||||
|
||||
assert result["final_response"] == "done"
|
||||
session_db.update_token_counts.assert_called_once()
|
||||
assert session_db.update_token_counts.call_args.args[0] == "cron-session"
|
||||
Reference in New Issue
Block a user