From d8e8fd81243c957e30e310282ca597a0540a5fc0 Mon Sep 17 00:00:00 2001
From: iomgaa <iomgaaycz@gmail.com>
Date: Thu, 30 Jul 2026 10:15:52 -0400
Subject: [PATCH] refactor: route TPM reservation and settlement through the
 derived value

Five call sites (QuotaGate entry, RetryMW/EmbeddingClient success and
transient-failure settlement) now read effective_est_tokens() instead of
est_tokens. Success paths gain an unavailable branch that keeps delta at
zero once usage frames may be missing; it has no producer yet, so
behaviour is unchanged while the tpm>0 => est_tokens>0 gate still holds.
---
 src/polygateway/embedding.py            |  9 +++++++--
 src/polygateway/middleware/ratelimit.py |  2 +-
 src/polygateway/middleware/retry.py     | 10 ++++++++--
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/polygateway/embedding.py b/src/polygateway/embedding.py
index d9512be..1bc7fd1 100644
--- a/src/polygateway/embedding.py
+++ b/src/polygateway/embedding.py
@@ -268,7 +268,11 @@ class EmbeddingClient:
                     source_name=source.name,
                     operation="embedding",
                 )
-            actual = result.prompt_tokens
+            if result.usage_source == "unavailable":
+                # 与 RetryMW 同口径: 用量不可得时按入场预扣量结算(设计 §3.2 #9)
+                actual = source.effective_est_tokens()
+            else:
+                actual = result.prompt_tokens
             await self._record_quietly(self._breaker.record_success(entry))
             await self._record_quietly(self._quota.mark_progress())
             latency_ms = int((self._now() - started) * 1000)
@@ -291,7 +295,8 @@ class EmbeddingClient:
             reasons[source.name] = reason
             await self._record_quietly(self._breaker.record_failure(entry, reason, dead))
             if not dead:
-                actual = source.est_tokens  # 保守: 失败请求可能已被网关计费(CHS 同款)
+                # 保守: 失败请求可能已被网关计费(CHS 同款);与入场预扣同源取值
+                actual = source.effective_est_tokens()
             await self._emit(batch, source, call_id, started, session_id, parent_call_id, error=exc)
             return _FailedBatch(exc, immediate=dead)
         finally:
diff --git a/src/polygateway/middleware/ratelimit.py b/src/polygateway/middleware/ratelimit.py
index f9ee398..7627ecd 100644
--- a/src/polygateway/middleware/ratelimit.py
+++ b/src/polygateway/middleware/ratelimit.py
@@ -23,7 +23,7 @@ class QuotaGate:
 
     async def try_acquire(self, source: SourceConfig) -> Permit | None:
         try:
-            return await self._limiter.try_acquire(source.name, source.est_tokens)
+            return await self._limiter.try_acquire(source.name, source.effective_est_tokens())
         except GovernanceBackendError:
             raise
         except Exception as exc:
diff --git a/src/polygateway/middleware/retry.py b/src/polygateway/middleware/retry.py
index 1c03eea..bb9e648 100644
--- a/src/polygateway/middleware/retry.py
+++ b/src/polygateway/middleware/retry.py
@@ -335,7 +335,12 @@ class RetryMW:
                 overlay=request.overlay,
                 call_id=call_id,
             )
-            actual = result.prompt_tokens + result.completion_tokens
+            if result.usage_source == "unavailable":
+                # 用量不可得时按入场预扣量结算(delta==0),否则押金会被整笔退回,
+                # 对"从不返回 usage 帧"的源等于 TPM 闸失效(设计 §3.2 #9)
+                actual = source.effective_est_tokens()
+            else:
+                actual = result.prompt_tokens + result.completion_tokens
             await self._record_quietly(self._breaker.record_success(entry))
             await self._record_quietly(self._quota.mark_progress())
             self._feed_outcome(source.name, ok=True)
@@ -367,7 +372,8 @@ class RetryMW:
                 self._pacer.on_backpressure(source.name)
             await self._record_quietly(self._breaker.record_failure(entry, reason, dead))
             if not dead:
-                actual = source.est_tokens  # 保守: 失败请求可能已被网关计费(CHS 同款)
+                # 保守: 失败请求可能已被网关计费(CHS 同款);与入场预扣同源取值
+                actual = source.effective_est_tokens()
             await self._emit(request, source, call_id, started, error=exc)
             return _Failed(exc, immediate=dead)
         finally: