fix(state-sync): mirror WebUI cache counters
This commit is contained in:
parent
1b58e4df47
commit
059a06b70d
4 changed files with 22 additions and 3 deletions
|
|
@ -101,6 +101,8 @@ def _sync_session_title_to_insights(session) -> None:
|
|||
title=session.title,
|
||||
message_count=len(messages),
|
||||
profile=getattr(session, "profile", None),
|
||||
cache_read_tokens=getattr(session, "cache_read_tokens", None) or 0,
|
||||
cache_write_tokens=getattr(session, "cache_write_tokens", None) or 0,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to update session title in state.db", exc_info=True)
|
||||
|
|
@ -20690,6 +20692,8 @@ def _handle_chat_sync(handler, body):
|
|||
model=s.model,
|
||||
title=s.title,
|
||||
message_count=len(s.messages),
|
||||
cache_read_tokens=s.cache_read_tokens or 0,
|
||||
cache_write_tokens=s.cache_write_tokens or 0,
|
||||
# #2762 / #2827 parity with api/streaming.py:5078: pass the
|
||||
# session's profile explicitly so a future refactor that
|
||||
# backgrounds this handler doesn't silently leak writes to
|
||||
|
|
|
|||
|
|
@ -134,7 +134,8 @@ def sync_session_start(session_id: str, model=None, profile: Optional[str] = Non
|
|||
|
||||
def sync_session_usage(session_id: str, input_tokens: int=0, output_tokens: int=0,
|
||||
estimated_cost=None, model=None, title: Optional[str] = None,
|
||||
message_count: Optional[int] = None, profile: Optional[str] = None) -> None:
|
||||
message_count: Optional[int] = None, profile: Optional[str] = None,
|
||||
cache_read_tokens: int = 0, cache_write_tokens: int = 0) -> None:
|
||||
"""Update token usage and title for a WebUI session in state.db.
|
||||
Called after each turn completes. Uses absolute=True to set totals
|
||||
(the WebUI Session already accumulates across turns).
|
||||
|
|
@ -153,11 +154,17 @@ def sync_session_usage(session_id: str, input_tokens: int=0, output_tokens: int=
|
|||
try:
|
||||
# Ensure session exists first (idempotent)
|
||||
db.ensure_session(session_id=session_id, source='webui', model=model)
|
||||
# Set absolute token counts
|
||||
# Set absolute token counts. WebUI's sidecar already accumulates
|
||||
# input/output/cache totals across turns, so mirror the same absolute
|
||||
# values into state.db. Omitting cache counters makes insights/reporting
|
||||
# show false 0% hit rates even when the live stream and sidecar saw warm
|
||||
# prefix reads.
|
||||
db.update_token_counts(
|
||||
session_id=session_id,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
estimated_cost_usd=estimated_cost,
|
||||
model=model,
|
||||
absolute=True,
|
||||
|
|
|
|||
|
|
@ -9287,6 +9287,8 @@ def _run_agent_streaming(
|
|||
model=model,
|
||||
title=s.title,
|
||||
message_count=len(s.messages),
|
||||
cache_read_tokens=s.cache_read_tokens or 0,
|
||||
cache_write_tokens=s.cache_write_tokens or 0,
|
||||
# #2762: pass the session's profile explicitly so the
|
||||
# background-thread state.db lookup doesn't fall
|
||||
# through to the process-global active profile and
|
||||
|
|
|
|||
|
|
@ -82,12 +82,14 @@ def _read_session(db_path: Path, session_id: str):
|
|||
if not db_path.exists():
|
||||
return None
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
try:
|
||||
# Real state.db schema (see api/state_sync.py + hermes_cli StateDB):
|
||||
# `sessions` table has `id` as PRIMARY KEY (not session_id). Use real
|
||||
# column names so the test queries the actual schema.
|
||||
cur = conn.execute(
|
||||
"SELECT id AS session_id, title, input_tokens, output_tokens "
|
||||
"SELECT id AS session_id, title, input_tokens, output_tokens, "
|
||||
"cache_read_tokens, cache_write_tokens "
|
||||
"FROM sessions WHERE id = ?",
|
||||
(session_id,),
|
||||
)
|
||||
|
|
@ -213,6 +215,8 @@ def test_sync_session_usage_writes_only_to_named_profile(two_profile_homes):
|
|||
title='2762 regression test',
|
||||
message_count=3,
|
||||
profile='maiko',
|
||||
cache_read_tokens=1200,
|
||||
cache_write_tokens=300,
|
||||
)
|
||||
|
||||
maiko_row = _read_session(two_profile_homes['maiko'] / 'state.db', '2762-regression')
|
||||
|
|
@ -220,6 +224,8 @@ def test_sync_session_usage_writes_only_to_named_profile(two_profile_homes):
|
|||
|
||||
assert maiko_row is not None, \
|
||||
"sync_session_usage(profile='maiko') did not write to maiko's state.db"
|
||||
assert maiko_row["cache_read_tokens"] == 1200
|
||||
assert maiko_row["cache_write_tokens"] == 300
|
||||
assert hiyuki_row is None, \
|
||||
"sync_session_usage(profile='maiko') leaked into hiyuki's state.db — #2762 regression"
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue