diff --git a/api/routes.py b/api/routes.py index 8bb31f1..baea175 100644 --- a/api/routes.py +++ b/api/routes.py @@ -1094,6 +1094,7 @@ def _handle_chat_sync(handler, body): estimated_cost=s.estimated_cost, model=s.model, title=s.title, + message_count=len(s.messages), ) except Exception: pass diff --git a/api/state_sync.py b/api/state_sync.py index 244e06f..9a80b0a 100644 --- a/api/state_sync.py +++ b/api/state_sync.py @@ -66,7 +66,8 @@ def sync_session_start(session_id: str, model=None) -> None: def sync_session_usage(session_id: str, input_tokens: int=0, output_tokens: int=0, - estimated_cost=None, model=None, title: str=None) -> None: + estimated_cost=None, model=None, title: str=None, + message_count: int=None) -> 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). @@ -92,6 +93,17 @@ def sync_session_usage(session_id: str, input_tokens: int=0, output_tokens: int= db.set_session_title(session_id, title) except Exception: pass + # Update message count + if message_count is not None: + try: + def _set_msg_count(conn): + conn.execute( + "UPDATE sessions SET message_count = ? WHERE id = ?", + (message_count, session_id), + ) + db._execute_write(_set_msg_count) + except Exception: + pass except Exception: pass # never crash the WebUI for sync failures finally: diff --git a/api/streaming.py b/api/streaming.py index b8bd051..6e92ccd 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -366,6 +366,7 @@ def _run_agent_streaming(session_id, msg_text, model, workspace, stream_id, atta estimated_cost=s.estimated_cost, model=model, title=s.title, + message_count=len(s.messages), ) except Exception: pass # never crash the stream for sync failures