From 1ff83bbfe0adf701000a6c8c578b5ed01e46241d Mon Sep 17 00:00:00 2001 From: iomgaa Date: Thu, 10 Sep 2026 00:48:07 -0400 Subject: [PATCH] fix: settle cancelled attempts against the source estimate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 取消发生在"端口已开始、结算尚未确定"时,原先按 settle(0) 把入场预扣整笔 退还,等于把可能已被上游计费的用量退回闸里;启用调用期限后库自身会常规性 触发该路径,故先修记账再启用。 改动只落在取消路径的取值上(设计 §6.3 矩阵 S3/S5/S7): - actual 初值保持 0,另设函数内局部阶段变量 settlement_known(不进任何签名); - 成功路径算出用量后置位,真实 usage 恰为 0 同样算"已知",不被取消覆写; - 已处理领域失败分支把结算决定前移到其第一个 await 之前(同级 except CancelledError 接不住本块 await 上的取消,它直穿 finally), 故 SourceDead 的既有 0 在取消下被保住,瞬时失败仍是 est,值与 1.3.5 逐字相同; - 未被四分类接住的异常不经上述分支,仍按 0 退全款(S8,本版不扩大语义); - OCR 的 0 token 是事实而非未知,settle(0) 不变,只补注释。 取消窗口一律用真实 asyncio.Event 钉死(不再 sleep 撞窗口),并加防越界回归: RuntimeError 逃逸仍结 0、真实 usage 为 0 的成功仍结 0。 --- src/polygateway/embedding.py | 16 ++- src/polygateway/middleware/retry.py | 18 ++- src/polygateway/ocr.py | 2 + .../test_redis_cross_connection.py | 30 ++++- tests/unit/test_embedding.py | 17 +++ tests/unit/test_ocr_client.py | 17 +++ tests/unit/test_retry.py | 121 +++++++++++++++++- 7 files changed, 213 insertions(+), 8 deletions(-) diff --git a/src/polygateway/embedding.py b/src/polygateway/embedding.py index 7b3a728..e191c2c 100644 --- a/src/polygateway/embedding.py +++ b/src/polygateway/embedding.py @@ -343,6 +343,9 @@ class EmbeddingClient: call_id = str(uuid.uuid4()) started = self._now() actual = 0 + # 局部阶段变量(与 RetryMW 同口径): 该刻库是否已算出确定结算。只服务于取消 + # 分支的兜底取值,不进任何签名; 未分类异常逃逸时仍逐字走旧的全额退还。 + settlement_known = False # 登记在 transport 调用**之前**(同 RetryMW): 失败与取消的尝试也真的发出去了 context.register_attempt() try: @@ -358,6 +361,8 @@ class EmbeddingClient: actual = source.effective_est_tokens() else: actual = result.prompt_tokens + # 真实 usage 恰为 0 也是已知事实, 后续取消不得改写成 est + settlement_known = True await self._record_quietly(self._breaker.record_success(entry)) await self._record_quietly(self._quota.mark_progress()) latency_ms = int((self._now() - started) * 1000) @@ -375,6 +380,7 @@ class EmbeddingClient: ) return _BatchOutcome(result, source, call_id, latency_ms) except (RequestRejectedError, ResultInvalidError) as exc: + actual, settlement_known = 0, True # 逐字保住 1.3.5 口径(本版不改这一族记账) await self._gate_on_terminal(exc, entry) await self._emit( batch, @@ -390,6 +396,9 @@ class EmbeddingClient: ) raise except asyncio.CancelledError: + if not settlement_known: + # 端口已开始、结算未定: 保守保留预扣(设计 §6.3 S7) + actual = source.effective_est_tokens() if entry.is_probe: await self._record_quietly(self._breaker.release_probe(entry)) await self._emit( @@ -409,10 +418,11 @@ class EmbeddingClient: dead = isinstance(exc, SourceDeadError) reason = _failure_reason(exc) reasons[source.name] = reason + # 结算决定定死在本分支第一个 await 之前: 同级 except CancelledError 接不住 + # 落在本块 await 上的取消,它直穿 finally。值与 1.3.5 逐字相同,只是算得更早。 + actual = 0 if dead else source.effective_est_tokens() + settlement_known = True await self._record_quietly(self._breaker.record_failure(entry, reason, dead)) - if not dead: - # 保守: 失败请求可能已被网关计费(CHS 同款);与入场预扣同源取值 - actual = source.effective_est_tokens() await self._emit( batch, source, diff --git a/src/polygateway/middleware/retry.py b/src/polygateway/middleware/retry.py index 7f66cd5..9e24d96 100644 --- a/src/polygateway/middleware/retry.py +++ b/src/polygateway/middleware/retry.py @@ -279,6 +279,9 @@ class RetryMW: call_id = str(uuid.uuid4()) started = self._now() actual = 0 + # 局部阶段变量: 该刻库是否已算出**确定**结算。只服务于取消分支的兜底取值, + # 不进任何签名/配置/遥测; 未分类异常逃逸时它无人读取, 故仍逐字走旧的全额退还。 + settlement_known = False # 登记在 transport 调用**之前**(1.3.5 设计 §4): 失败与取消的尝试同样 # "真的打出去了",挪到成功之后会让诊断最需要看见的那几次从计数里消失。 # 上下文为 None = 库内现场构造的请求,跳过而不是报错 @@ -301,6 +304,8 @@ class RetryMW: actual = source.effective_est_tokens() else: actual = result.prompt_tokens + result.completion_tokens + # 真实 usage 恰为 0 也是**已知事实**, 后续取消不得把它改写成 est + settlement_known = True await self._record_quietly(self._breaker.record_success(entry)) await self._record_quietly(self._quota.mark_progress()) self._feed_outcome(source.name, ok=True) @@ -309,15 +314,20 @@ class RetryMW: await self._emit(request, source, call_id, started, response=response) return response except RequestRejectedError as exc: + actual, settlement_known = 0, True # 逐字保住 1.3.5: 坏请求全额退还 await self._on_rejected(exc, source, entry) await self._emit(request, source, call_id, started, error=exc) raise except ResultInvalidError as exc: + actual, settlement_known = 0, True # 同上, 本版不改这一族记账口径 # 坏结果 ≠ 坏服务: 熔断记成功但不计窗口样本,亦不喂健康分(M2.5 §3.1) await self._record_quietly(self._breaker.record_success(entry, count_attempt=False)) await self._emit(request, source, call_id, started, error=exc) raise except asyncio.CancelledError: + if not settlement_known: + # 端口已开始、结算未定: 保守保留预扣(宁多扣不凭空退款, 见设计 §6.3 S3) + actual = source.effective_est_tokens() if entry.is_probe: await self._record_quietly(self._breaker.release_probe(entry)) await self._emit(request, source, call_id, started, error="cancelled") @@ -330,10 +340,12 @@ class RetryMW: self._feed_outcome(source.name, ok=False) if reason == "rate_limited": self._pacer.on_backpressure(source.name) + # 结算决定必须在本分支**第一个 await 之前**定死: 同级的 except CancelledError + # 接不住落在本块 await 上的取消,它会直穿 finally——那一刻 actual 是什么就结什么。 + # 值与 1.3.5 逐字相同(dead 全额退、瞬时保留预扣),只是算得更早。 + actual = 0 if dead else source.effective_est_tokens() + settlement_known = True await self._record_quietly(self._breaker.record_failure(entry, reason, dead)) - if not dead: - # 保守: 失败请求可能已被网关计费(CHS 同款);与入场预扣同源取值 - actual = source.effective_est_tokens() await self._emit(request, source, call_id, started, error=exc) return _Failed(exc, immediate=dead) finally: diff --git a/src/polygateway/ocr.py b/src/polygateway/ocr.py index 501ae18..2c282d7 100644 --- a/src/polygateway/ocr.py +++ b/src/polygateway/ocr.py @@ -446,6 +446,8 @@ class OcrClient: ) return _FailedAttempt(exc, immediate=dead) finally: + # OCR 的 0 token 是**事实**而非"未知"(设计 §6.3 S6): 故取消也恰恰结 0, + # 不引入 chat/embedding 那套 settlement_known 兜底。 await settle_and_release(permit, 0) async def _invoke( diff --git a/tests/integration/test_redis_cross_connection.py b/tests/integration/test_redis_cross_connection.py index 4b37cb6..5405154 100644 --- a/tests/integration/test_redis_cross_connection.py +++ b/tests/integration/test_redis_cross_connection.py @@ -72,10 +72,13 @@ class ScriptedTransport: def __init__(self, hang: bool = False): self.hang = hang self.calls: list[str] = [] + # 取消用例的确定性窗口(同 test_retry FakeTransport): 进入挂起即置位 + self.entered = asyncio.Event() async def complete(self, *, messages, source, stream, overlay, call_id, reasoning_effort): self.calls.append(source.name) if self.hang: + self.entered.set() await asyncio.Event().wait() return TransportResult( content="ok", @@ -225,7 +228,32 @@ async def test_cancel_in_flight_releases_lease(clients): assert (await limiter.source_stats("s1")).inflight == 0 -# —— 掉线方向(fail-closed 集成证据)—— +async def test_cancel_in_flight_keeps_the_reservation(clients): + """1.3.6 §6.3 S3 在**真实 Redis** 上: 端口已开始、用量未知 → 保留 est 预扣。 + + 内存后端与 Lua 后端的 `settle(delta)` 算术必须同口径——取消时凭空退款 + 在分布式部署下就是几个 worker 一起击穿 TPM 闸。本用例不改 Lua、不改契约套件。 + """ + a_cli, _ = clients + scope = f"t{uuid4().hex[:8]}" + sources = [make_source(max_concurrency=1, tpm=1000, est_tokens=400)] + limiter = _limiter(a_cli, scope, sources, GlobalLimits(0, 0, 0)) + transport = ScriptedTransport(hang=True) + client = _client( + scope, + sources, + limiter, + RedisGate(config=_CFG, redis=a_cli, scope=scope), + transport, + ) + task = asyncio.create_task(client.chat([{"role": "user", "content": "hi"}])) + await transport.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + stats = await limiter.source_stats("s1") + assert stats.tpm_used == 400 # 预扣保留, 不回退到 0 + assert stats.inflight == 0 async def test_redis_down_admission_fails_closed(): diff --git a/tests/unit/test_embedding.py b/tests/unit/test_embedding.py index f6ad5a3..624d1d6 100644 --- a/tests/unit/test_embedding.py +++ b/tests/unit/test_embedding.py @@ -197,6 +197,8 @@ class ScriptedEmbedTransport: def __init__(self, script): self.script = list(script) self.calls = [] + # 取消用例的确定性窗口: 进入 hang 分支即置位, 不用 sleep 撞窗口 + self.entered = asyncio.Event() async def embed(self, *, texts, source, call_id): self.calls.append((source.name, list(texts), call_id)) @@ -204,6 +206,7 @@ class ScriptedEmbedTransport: if isinstance(action, Exception): raise action if action == "hang": + self.entered.set() await asyncio.Event().wait() if action == "ok": return _vec_for(texts) @@ -365,6 +368,20 @@ class TestEmbedGovernance: await task assert (await limiter.source_stats("e1")).inflight == 0 + async def test_cancel_in_flight_keeps_the_reservation(self): + """S7(与 chat 同口径): transport 在途被取消 → 用量未知 → 保留预扣而非退成 0。""" + transport = ScriptedEmbedTransport(["hang"]) + client, limiter = _embed_client( + [_src(max_concurrency=1, tpm=1000, est_tokens=400)], [], transport=transport + ) + task = asyncio.create_task(client.embed(["a"])) + await transport.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + stats = await limiter.source_stats("e1") + assert stats.tpm_used == 400 and stats.inflight == 0 + async def test_single_timeout_does_not_exhaust_stall_budget(self): """issue #8: 一次耗满 timeout 的尝试不得吃掉 stall 预算而使重试失效。 diff --git a/tests/unit/test_ocr_client.py b/tests/unit/test_ocr_client.py index d4f03d0..438b0f9 100644 --- a/tests/unit/test_ocr_client.py +++ b/tests/unit/test_ocr_client.py @@ -60,6 +60,8 @@ class ScriptedOcrTransport: def __init__(self, script): self.script = list(script) self.calls = [] + # 取消用例的确定性窗口: 进入 hang 分支即置位, 不用 sleep 撞窗口 + self.entered = asyncio.Event() async def _next(self, method, source, call_id): self.calls.append((method, source.name, call_id)) @@ -67,6 +69,7 @@ class ScriptedOcrTransport: if isinstance(action, Exception): raise action if action == "hang": + self.entered.set() await asyncio.Event().wait() return _TEXT_OK if action == "text" else _LAYOUT_OK @@ -395,6 +398,20 @@ class TestCancellation: stats = await limiter.source_stats("m1") assert stats.inflight == 0 # permit 在 finally 释放 + async def test_cancel_in_flight_still_settles_zero(self): + """S6: OCR 的 0 token 是**事实**而非"未知", 取消也不得改成按 est 结算。""" + transport = ScriptedOcrTransport(["hang"]) + client, limiter, _ = _client( + [_src(max_concurrency=1, tpm=1000, est_tokens=400)], [], transport=transport + ) + task = asyncio.create_task(client.recognize_text(b"jpg")) + await transport.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + stats = await limiter.source_stats("m1") + assert stats.tpm_used == 0 and stats.inflight == 0 + class TestCheckHealth: class _HealthTransport(ScriptedOcrTransport): diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py index 5a91b7a..efc7b81 100644 --- a/tests/unit/test_retry.py +++ b/tests/unit/test_retry.py @@ -76,6 +76,8 @@ class FakeTransport: self.script = list(script) self.calls = [] self.efforts = [] + # 取消用例的确定性窗口: 进入 hang 分支即置位, 用例据此取消而非 sleep 猜时长 + self.entered = asyncio.Event() async def complete(self, *, messages, source, stream, overlay, call_id, reasoning_effort): self.calls.append((source.name, call_id)) @@ -84,10 +86,37 @@ class FakeTransport: if isinstance(action, Exception): raise action if action == "hang": + self.entered.set() await asyncio.Event().wait() return action +class HangingGate(InMemoryGate): + """在指定记账写回处永久挂起的门控: 把"取消落在某个 await 上"变成确定性事件。 + + 只覆盖 `record_success` / `record_failure` 两个写回点, 其余行为沿用真实内存实现。 + """ + + def __init__(self, *, hang_on, **kwargs): + super().__init__(**kwargs) + self._hang_on = hang_on + self.entered = asyncio.Event() + + async def _hang(self): + self.entered.set() + await asyncio.Event().wait() + + async def record_success(self, entry, *, count_attempt=True): + if self._hang_on == "success": + await self._hang() + return await super().record_success(entry, count_attempt=count_attempt) + + async def record_failure(self, entry, reason, force_open): + if self._hang_on == "failure": + await self._hang() + return await super().record_failure(entry, reason, force_open) + + class FakeSleep: """记录退避时长,立即返回(不真等)。""" @@ -109,6 +138,7 @@ def _harness( rng=lambda: 0.0, selector=None, pacer=None, + gate=None, ): clock = clock or FakeClock() limiter = InMemoryLimiter( @@ -118,7 +148,7 @@ def _harness( lease_ttl_s=100.0, now=clock, ) - gate = InMemoryGate(config=_BREAKER, now=clock) + gate = gate if gate is not None else InMemoryGate(config=_BREAKER, now=clock) transport = FakeTransport(script) sleep = FakeSleep() mw = RetryMW( @@ -467,6 +497,95 @@ class TestScopeUnavailable: assert resp.content == "ok" and released["done"] +class TestCancellationSettlement: + """1.3.6 §6.3 结算矩阵: 取消时 `settle()` 的取值只由"该刻库知道什么"决定。 + + 取消窗口一律用真实 `asyncio.Event` 钉死(不用 sleep 撞窗口), 否则红绿都不可信。 + """ + + async def test_cancel_in_flight_keeps_the_reservation(self): + """S3: transport 在途被取消 → 端口已开始、用量未知 → 保留预扣(不凭空退款)。""" + mw, limiter, _, transport, *_ = _harness( + [_src("a", max_concurrency=1, tpm=1000, est_tokens=400)], ["hang"] + ) + task = asyncio.ensure_future(mw(_REQ)) + await transport.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + stats = await limiter.source_stats("a") + assert stats.tpm_used == 400 # est 保留, 而非退成 0 + assert stats.inflight == 0 + + async def test_cancel_after_usage_known_keeps_real_usage(self): + """S4: 真实 usage 已算出后被取消 → 结算仍是真实值, 不被 est 覆写。""" + clock = FakeClock() + gate = HangingGate(hang_on="success", config=_BREAKER, now=clock) + mw, limiter, *_ = _harness( + [_src("a", max_concurrency=1, tpm=1000, est_tokens=400)], + [_ok()], + clock=clock, + gate=gate, + ) + task = asyncio.ensure_future(mw(_REQ)) + await gate.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert (await limiter.source_stats("a")).tpm_used == 15 # 10+5 实测 + + async def test_cancel_in_dead_failure_branch_keeps_full_refund(self): + """S5-dead: 源已判死时的既有 `0` 不得因取消退化成 est(不得继续占额度)。""" + clock = FakeClock() + gate = HangingGate(hang_on="failure", config=_BREAKER, now=clock) + mw, limiter, *_ = _harness( + [_src("a", max_concurrency=1, tpm=1000, est_tokens=400)], + [SourceDeadError("401", source_name="a", status_code=401)], + clock=clock, + gate=gate, + ) + task = asyncio.ensure_future(mw(_REQ)) + await gate.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert (await limiter.source_stats("a")).tpm_used == 0 + + async def test_cancel_in_transient_failure_branch_keeps_the_reservation(self): + """S5-transient: 瞬时失败的结算决定在首个 await 之前定死, 取消拿到同一个 est。""" + clock = FakeClock() + gate = HangingGate(hang_on="failure", config=_BREAKER, now=clock) + mw, limiter, *_ = _harness( + [_src("a", max_concurrency=1, tpm=1000, est_tokens=400)], + [TransientError("boom", source_name="a", status_code=500)], + clock=clock, + gate=gate, + ) + task = asyncio.ensure_future(mw(_REQ)) + await gate.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert (await limiter.source_stats("a")).tpm_used == 400 + + async def test_unclassified_exception_still_refunds_in_full(self): + """S8 防越界: 未分类异常(无 except 接住)仍逐字走 1.3.5 的全额退还。""" + mw, limiter, *_ = _harness( + [_src("a", max_concurrency=1, tpm=1000, est_tokens=400)], [RuntimeError("boom")] + ) + with pytest.raises(RuntimeError): + await mw(_REQ) + stats = await limiter.source_stats("a") + assert stats.tpm_used == 0 and stats.inflight == 0 + + async def test_real_zero_usage_success_settles_zero(self): + """防越界: 真实 usage 恰为 0 是**已知事实**, 不得被当成"未知"改按 est 结算。""" + zero = dataclasses.replace(_ok(), prompt_tokens=0, completion_tokens=0) + mw, limiter, *_ = _harness([_src("a", tpm=1000, est_tokens=400)], [zero]) + await mw(_REQ) + assert (await limiter.source_stats("a")).tpm_used == 0 + + class TestCancellation: async def test_cancel_mid_flight_releases_permit(self): mw, limiter, _, _, _, _ = _harness([_src("a", max_concurrency=1)], ["hang"])