From 7d69b80012ea19744efc7b59314bd90a346029b5 Mon Sep 17 00:00:00 2001
From: Frank Song <franksong2702@gmail.com>
Date: Thu, 25 Jun 2026 22:11:46 +0800
Subject: [PATCH] Avoid repeated session loads in live usage metering

---
 CHANGELOG.md                                |  4 +++
 api/streaming.py                            | 40 ++++++++++++++++++---
 tests/test_streaming_live_usage_estimate.py | 30 +++++++++++++++-
 3 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5c6ddf2c..3440c7b0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,10 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- **Live metering during active streams no longer reloads the session on every usage tick.** The streaming worker now reuses its current session object for `_live_usage_snapshot()` and only falls back to `get_session()` once before the worker has loaded that object. This removes another avoidable global-session-lock touch point while tokens and tool events are streaming, narrowing the lock-contention half of #4918 without changing usage payloads or active-stream state.
+
 ## [v0.51.654] — 2026-06-25 — Release XJ (session saves block reads less during streaming)
 
 ### Changed
diff --git a/api/streaming.py b/api/streaming.py
index 02eb64bb..14b61a9b 100644
--- a/api/streaming.py
+++ b/api/streaming.py
@@ -4923,6 +4923,31 @@ def live_usage_prompt_estimate_after_tool_delta(
     }
 
 
+def _live_usage_session_snapshot(session_id, current_session, cache_ref, *, loader=get_session):
+    """Return a session object for hot live-metering paths without repeated loads."""
+    if current_session is not None:
+        try:
+            cache_ref[0] = current_session
+        except Exception:
+            pass
+        return current_session
+    try:
+        cached = cache_ref[0]
+    except Exception:
+        cached = None
+    if cached is not None:
+        return cached
+    try:
+        loaded = loader(session_id)
+    except Exception:
+        return None
+    try:
+        cache_ref[0] = loaded
+    except Exception:
+        pass
+    return loaded
+
+
 def _tool_result_snippet(raw, limit: int = _TOOL_RESULT_SNIPPET_MAX) -> str:
     """Extract a bounded result preview from a stored tool message payload."""
     if limit <= 0:
@@ -5744,6 +5769,14 @@ def _run_agent_streaming(
     # (model, base_url, provider) within one stream, so resolve it at most
     # once. Sentinel: None=not computed, 0=not applicable/failed, >0=real cap.
     _real_ctx_cache = [None]
+    _live_usage_session_cache = [None]
+
+    def _current_live_usage_session():
+        return _live_usage_session_snapshot(
+            session_id,
+            s,
+            _live_usage_session_cache,
+        )
 
     def _seed_live_prompt_estimate() -> int:
         """Capture the latest exact prompt size before adding live tool deltas."""
@@ -5760,7 +5793,7 @@ def _run_agent_streaming(
                 _base = 0
         if not _base:
             try:
-                _session_obj = get_session(session_id)
+                _session_obj = _current_live_usage_session()
                 _base = getattr(_session_obj, 'last_prompt_tokens', 0) or 0
             except Exception:
                 _base = 0
@@ -5802,10 +5835,7 @@ def _run_agent_streaming(
             'threshold_tokens': 0,
             'last_prompt_tokens': 0,
         }
-        try:
-            _session_obj = get_session(session_id)
-        except Exception:
-            _session_obj = None
+        _session_obj = _current_live_usage_session()
 
         _agent = agent
         if _agent is not None:
diff --git a/tests/test_streaming_live_usage_estimate.py b/tests/test_streaming_live_usage_estimate.py
index 9f1edb04..689474db 100644
--- a/tests/test_streaming_live_usage_estimate.py
+++ b/tests/test_streaming_live_usage_estimate.py
@@ -1,4 +1,7 @@
-from api.streaming import live_usage_prompt_estimate_after_tool_delta
+from api.streaming import (
+    _live_usage_session_snapshot,
+    live_usage_prompt_estimate_after_tool_delta,
+)
 
 
 def test_live_usage_estimate_caps_tool_delta_against_previous_prompt():
@@ -48,3 +51,28 @@ def test_live_usage_estimate_caps_cumulative_tool_delta_per_turn():
     assert usage["turn_tool_prompt_tokens"] == 24_000
     assert usage["last_prompt_tokens"] == base_prompt_tokens + 24_000
     assert usage["last_prompt_tokens"] < base_prompt_tokens + (20 * 12_000)
+
+
+def test_live_usage_session_snapshot_prefers_current_stream_session():
+    current = object()
+    cache = [None]
+
+    def fail_loader(_sid):
+        raise AssertionError("hot live usage path should not reload the current session")
+
+    assert _live_usage_session_snapshot("sid", current, cache, loader=fail_loader) is current
+    assert cache[0] is current
+
+
+def test_live_usage_session_snapshot_falls_back_once_then_uses_cache():
+    loaded = object()
+    calls = []
+    cache = [None]
+
+    def loader(sid):
+        calls.append(sid)
+        return loaded
+
+    assert _live_usage_session_snapshot("sid", None, cache, loader=loader) is loaded
+    assert _live_usage_session_snapshot("sid", None, cache, loader=loader) is loaded
+    assert calls == ["sid"]