diff --git a/app/client_hangup.py b/app/client_hangup.py new file mode 100644 index 0000000..3ba26e8 --- /dev/null +++ b/app/client_hangup.py @@ -0,0 +1,41 @@ +"""`stream=false` 的聚合窗口：监听下游断连，并把没听完的那一枪取消掉。""" +from __future__ import annotations + +import asyncio + + +class ClientHungUp(Exception): + """客户端在响应成形之前挂断了。不是错误，是「没人听了」。""" + + +async def _listen_for_hangup(request): + """等 `http.disconnect`；多余的 `http.request` 残留消息不算断连。""" + while True: + message = await request.receive() + if message.get("type") == "http.disconnect": + return + + +async def await_or_hangup(awaitable, request): + """等待一次上游调用；期间下游挂断就取消它。返回原调用的结果，异常原样抛出。 + + `request is None` 时退化成普通的 `await`。取消必须等收尾跑完再交出去：httpx 的 + `async with` 在 `finally` 里才 `aclose()` 响应，否则上游连接会留在半关状态。 + """ + if request is None: + return await awaitable + work = asyncio.ensure_future(awaitable) + watch = asyncio.ensure_future(_listen_for_hangup(request)) + try: + done, _ = await asyncio.wait((work, watch), return_when=asyncio.FIRST_COMPLETED) + if work in done: + return work.result() # 上游先回来：这一次结果是真的完整，照旧交给服务端投递 + work.cancel() + await asyncio.wait((work,)) + if not work.cancelled(): + work.exception() # 取消途中自己报了错：也属于「没人听了」，但不留未取回的异常 + raise ClientHungUp + finally: + work.cancel() + watch.cancel() + await asyncio.gather(work, watch, return_exceptions=True) diff --git a/converter.py b/converter.py index 43e829a..3b5590f 100644 --- a/converter.py +++ b/converter.py @@ -43,7 +43,7 @@ import httpx from fastapi import FastAPI, Header, HTTPException, Request from fastapi.exception_handlers import http_exception_handler as _default_http_exception_handler -from fastapi.responses import JSONResponse, StreamingResponse +from fastapi.responses import JSONResponse, Response, StreamingResponse from starlette.concurrency import run_in_threadpool import uvicorn @@ -69,6 +69,7 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False, from app import trial_rewards from app import checkin as checkin_service, model_policy, travel from app.model_blocks import ModelBlocks +from app.client_hangup import ClientHungUp, await_or_hangup from app.observability import (AuditMiddleware, observe_recovery, observe_route, observe_usage, observe_attempt, observe_failure, observe_failure_seq) @@ -2415,8 +2416,13 @@ def attempt(routed, cred, headers, url): async def fetch(routed, cred, headers, url): return await _fetch_checked_chat(url, headers, routed, model_name, rid, cred, filter_retry=True) - collected = await _routed_fetch(payload, prepared, model_name, rid, t0, fetch, - body, cred, headers, url) + # 断连监听包在重放外层：换凭证的那几枪同样属于「还没给下游一个字节」的窗口 + try: + collected = await await_or_hangup( + _routed_fetch(payload, prepared, model_name, rid, t0, fetch, + body, cred, headers, url), request) + except ClientHungUp: + return _hungup_response(rid, model_name, t0) _log_finish(model_name, t0, collected, rid) if CONFIG.get("control_store") is not None: collected = {**collected, "model": model_name} @@ -2653,6 +2659,17 @@ def _upstream_failure(error, model_name, t0, rid): return status, raw +def _hungup_response(rid, model_name, t0): + """下游已经不听了：安静地给一个不成体的响应，不编造结果。 + + 204 只是「ASGI 调用必须交付一个响应」的形式（Starlette 对 204 不写 content-length）； + 审计由 `AuditMiddleware` 判定为 cancelled。 + """ + elapsed = time.time() - t0 if t0 else 0 + _log(f"[{rid}] ✂ 下游已断连，取消这次聚合 | {model_name} | {elapsed:.1f}s") + return Response(status_code=204) + + async def _fetch_checked_chat(url, headers, body, model_name, rid, cred=None, *, filter_retry=False): """统一聚合与校验；非流式纯审核拒绝最多压缩兜底一次，网络错误不重放。""" tool_attempt = 0 @@ -3170,25 +3187,28 @@ def attempt(routed, cred, headers, url): chat_body, cred, headers, url) return await _nonstream_adapted(url, headers, chat_body, model_name, t0, rid, cred, - payload=payload, canonical=prepared) + payload=payload, canonical=prepared, request=request) async def _nonstream_adapted(url, headers, body, model_name, t0, rid, cred, *, anthropic=False, - payload=None, canonical=None): + payload=None, canonical=None, request=None): converter = (AnthropicStreamConverter(model=model_name) if anthropic else ResponsesStreamConverter(model=model_name, parallel_tool_calls=body.get("parallel_tool_calls", True))) async def fetch(routed, cred, headers, url): return await _fetch_checked_chat(url, headers, routed, model_name, rid, cred, filter_retry=True) try: - collected = await _routed_fetch(payload, body if canonical is None else canonical, - model_name, rid, t0, fetch, body, cred, headers, url) + collected = await await_or_hangup( + _routed_fetch(payload, body if canonical is None else canonical, + model_name, rid, t0, fetch, body, cred, headers, url), request) for line in _chat_result_to_sse_lines(_completion_to_merged(collected)): converter.feed_line(_public_sse_line(line, model_name)) converter.finish() except (httpx.HTTPError, UpstreamResponseError) as error: status, raw = _upstream_failure(error, model_name, t0, rid) raise HTTPException(status_code=status, detail=_safe_err_raw(raw, status)) from None + except ClientHungUp: + return _hungup_response(rid, model_name, t0) result = converter.get_nonstream_response() _log_finish(model_name, t0, collected, rid) return JSONResponse(content=result) @@ -3272,7 +3292,8 @@ async def create_message(request: Request, if not _client_wants_stream(payload): return await _nonstream_adapted(url, headers, chat_body, model_name, t0, rid, cred, - anthropic=True, payload=payload, canonical=prepared) + anthropic=True, payload=payload, canonical=prepared, + request=request) def attempt(routed, cred, headers, url): return _stream_anthropic(url, headers, routed, model_name, t0, rid, cred=cred) diff --git a/docs/advanced.md b/docs/advanced.md index 1e1edbf..74995b6 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -175,6 +175,7 @@ Credential domain / token issuer determine the product identity. Chat and refres | Upstream `service info not found` (code 11102) | That backend does not serve the model at all: avoid it for the `(backend, model)` pair, route the model to another backend, and return 404 when none has it. Half-open after 6 h, exponential backoff up to 24 h, cleared at once by one successful call; inspect via `GET /admin/model-blocks` | | Connection setup failure | Retry once on a fresh connection: `ConnectError` and `ConnectTimeout` fail before the first body byte, so the upstream holds nothing and replaying cannot double-bill | | Post-send disconnect, read timeout or protocol error | No network replay, avoiding duplicate billing; logs include exception type and elapsed time | +| Client hangs up before a non-streaming response is ready | The upstream call is cancelled and its concurrency slot returned at once; the request is audited as `cancelled`, never as a completed answer. Streaming already behaves this way | | Streaming request fails before the first byte | Reported with the real HTTP status, exactly like `stream=false`. A 200 carrying only an in-band `error` event is read by clients as an empty answer, so the session ends silently while the audit log records a success | | Credential failover (`--failover-max`) | Off by default. When enabled, a failure that happened before any byte reached the client is retried on another credential, up to N times, and is audited as `success` with a `failover_recovered` attempt marker. Only upstream HTTP rejections (401/403/429/502/503/504) and request bodies the upstream provably never started receiving (`ConnectError`/`ConnectTimeout`) qualify: content-filter rejections, 502s synthesized from an already-open stream, read timeouts and protocol errors are never replayed, and if no other credential is available the original status is surfaced — **Billing**: 401/403/429/503 and those transport failures happen at admission time and cannot be billed; a 502/504 may already have been processed and billed upstream, but its result never reached the client, so refusing to replay it recovers no credit — it only turns a paid-for attempt into a broken session. Such replays are tagged `上游可能已处理该请求` in the log for reconciliation | | Write-timeout replay (`--retry-write-timeout`) | Off by default. A write timeout proves the declared body was not fully sent, not that the upstream ignored the bytes it did receive, so it is excluded from both the connect retry and credential failover until explicitly enabled. Long cross-border sessions fail here more often than in the handshake, so operators who have confirmed their upstream does not bill partial bodies can turn this on; those replays carry the same `上游可能已处理该请求` log tag | diff --git a/docs/advanced.zh-CN.md b/docs/advanced.zh-CN.md index 9207c1d..6cd7cde 100644 --- a/docs/advanced.zh-CN.md +++ b/docs/advanced.zh-CN.md @@ -172,6 +172,7 @@ WebUI 可以直接上传文件；以下限制针对 `POST /admin/credentials` | 上游 `service info not found`（11102） | 该后端根本不服务这个模型：按 (后端, 模型) 避让，把该模型派给其他后端，全部后端都没有时返回 404。6 小时后半开放行重试，反复命中最长退避 24 小时，一次成功调用即刻解除；可用 `GET /admin/model-blocks` 查看 | | 建连失败 | `ConnectError` / `ConnectTimeout` 换新连接重放一次：两者都发生在写下第一个正文字节之前，上游手里什么都没有，重放不会重复计费 | | 发送后断连、读超时、协议错误 | 不做网络重放，避免重复计费；日志记录异常类型与耗时 | +| 非流式响应还没成形，客户端就挂断 | 立刻取消这次上游调用并归还并发名额，审计记为 `cancelled`，不会被记成一次已完成的回答；流式本来就是这一行为 | | 流式在第一个字节之前失败 | 按**真实状态码**返回，与 `stream=false` 同口径。只带一个流内 `error` 事件的 200 会被客户端读成「模型答了个空」，会话静默结束，审计里还记成一次成功 | | 换凭证重放（`--failover-max`） | 默认关闭。开启后，失败落在「一个字节都没发给下游」之前时换一个凭证重打，最多 N 次，审计记为 `success` 并留下 `failover_recovered` 尝试标记。只重放上游用 HTTP 状态码给出的拒绝（401/403/429/502/503/504）与确定没开始收正文的传输失败（建连失败/建连超时）：内容审核拒绝、上游已回 200 之后合成的 502、读超时与协议错误一律不重放；换不出其他凭证时如实回第一次的状态码。**计费口径**：401/403/429/503 与建连类失败都发生在受理阶段，不会扣费；502/504 有可能已被后端处理并计费，但那次结果对下游根本拿不到，不重放也退不回额度——只是把一次已经付费的请求换成一段断掉的会话。这类重放在日志里单独标注「上游可能已处理该请求」，便于按官方用量明细核对 | | 写超时重放（`--retry-write-timeout`） | 默认关闭。写超时只能证明声明的正文没发完，不能证明上游忽略了已经收到的那部分，因此它默认既不参与连接重试也不参与换凭证重放；显式打开后两层都生效。跨境长会话最容易撞的恰恰是 60s 写超时（实测某部署传输失败 100% 是它），确认自己的上游不会按半截正文计费再开；这类重放同样带「上游可能已处理该请求」日志标记 | diff --git a/tests/test_nonstream_disconnect.py b/tests/test_nonstream_disconnect.py new file mode 100644 index 0000000..000edd5 --- /dev/null +++ b/tests/test_nonstream_disconnect.py @@ -0,0 +1,315 @@ +#!/usr/bin/env python3 +"""`stream=false` 的聚合窗口监听下游断连：取消上游、归还名额、审计如实，三协议一致。 + +流式端点由 Starlette 的 `listen_for_disconnect` 兜住，非流式端点没有对应机制。这里钉住 +聚合路径的边界，并覆盖换凭证重放途中挂断、以及调用方自己取消外层任务两种情形。 + +运行：.venv/bin/python -B -m unittest -v tests/test_nonstream_disconnect.py +""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # 仓库根：允许直接运行本文件 + +import asyncio +import json +import unittest +from contextlib import contextmanager, suppress +from unittest.mock import patch + +import httpx +from fastapi import FastAPI + +import converter +from app.audit_store import AuditStore +from app.inbound_limits import ConcurrencyLimitMiddleware +from app.observability import AuditMiddleware +from tests import test_region_routing as fixtures + +HANG_MARKER = "synthetic-hang-up" +# 客户端已经不在了：唯一不可接受的是一个看起来正常的推理响应 +QUIET_STATUSES = (None, 204) +STEP = 2 +ENDPOINTS = fixtures.GENERATIONS +FAILOVER_LOG = "换凭证重放" + + +def error_body(message="synthetic rejection", code="rate_limit"): + return {"error": {"message": message, "type": "upstream_error", "code": code}} + + +@contextmanager +def allow_failover(times: int): + """打开换凭证重放开关（等价于 --failover-max N）。""" + with patch.dict(converter.CONFIG, {"failover_max": times}): + yield + + +class _HangingStream(httpx.AsyncByteStream): + """永不结束的上游 SSE：占住整个聚合窗口。 + + 两个独立标记：`read_cancelled` 只在 `__aiter__` 的 `finally` 置位，证明读取被取消； + `stream_closed` 由 httpx 的 `Response.aclose()` 调到底，证明连接真的还回去了。 + """ + + def __init__(self, read_cancelled, closed): + self.read_cancelled = read_cancelled + self.closed = closed + + async def __aiter__(self): + try: + await asyncio.Event().wait() # 第一段永远不来，等价于上游卡住 + yield b"" + finally: + self.read_cancelled.set() + + async def aclose(self): + self.closed.set() + await super().aclose() + + +class NonStreamDisconnectTests(fixtures.RegionRoutingTests): + """自己驱动 ASGI：需要一个能随时吐出 `http.disconnect` 的 receive，TestClient 给不了。""" + + def setUp(self): + super().setUp() + self.allowed_profiles = set(fixtures.PROFILES) + self.reset_upstream() + + def reset_upstream(self): + """每个协议子例各自从干净的上游状态起（subTest 共用同一次 setUp）。""" + self.upstream_seen = asyncio.Event() + self.read_cancelled = asyncio.Event() + self.stream_closed = asyncio.Event() + self.hang_attempts = 0 + self.hang_uids = [] + self.fail_over_first = False + + # --- 上游夹具：带标记的那一枪卡在首段之前；被点名时先让第一枪吃 429 --- + def handle_upstream(self, request): + if HANG_MARKER in request.content.decode("utf-8", "replace"): + self.hang_attempts += 1 + self.hang_uids.append(request.headers.get("x-user-id")) + if self.hang_attempts == 1 and self.fail_over_first: + return httpx.Response(429, json=error_body(), + headers={"content-type": "application/json"}) + self.upstream_seen.set() + # 用 stream= 而不是 content=：后者会把流再包一层，自定义 aclose 收不到关闭回调 + return httpx.Response(200, stream=_HangingStream(self.read_cancelled, + self.stream_closed), + headers={"content-type": "text/event-stream"}) + return super().handle_upstream(request) + + # --- 驱动 --- + def gated(self, limit, store=None): + """本用例私有的名额闸：`converter.app` 自带的那个跨用例复用，会把状态串到别的测试上。""" + application = converter.app + if store is not None: + application = FastAPI() + application.router.routes = list(converter.app.router.routes) + application.add_middleware(AuditMiddleware, {"audit_store": store}) + return ConcurrencyLimitMiddleware(application, {"max_concurrent": limit}) + + def scope(self, endpoint="chat/completions"): + path = "/v1/" + endpoint + return {"type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", "method": "POST", "scheme": "http", + "path": path, "raw_path": path.encode(), "query_string": b"", "root_path": "", + "server": ("testserver", 80), "client": ("127.0.0.1", 55555), + "headers": [(b"host", b"testserver"), (b"content-type", b"application/json")]} + + def receive(self, payload, hangup): + """一份正常请求体；之后把 `http.disconnect` 攥在手里，等这个「客户端」真的挂断。""" + pending = [{"type": "http.request", "body": json.dumps(payload).encode(), + "more_body": False}] + + async def receive(): + if pending: + return pending.pop(0) + await hangup.wait() + return {"type": "http.disconnect"} + return receive + + @staticmethod + def collect(sent): + async def send(message): + sent.append(message) + return send + + def started_statuses(self, sent): + return [m.get("status") for m in sent if m.get("type") == "http.response.start"] + + def hang_payload(self, endpoint="chat/completions", *, stream=False): + return self.payload(endpoint, text=HANG_MARKER, stream=stream) + + async def drain_records(self, store, expected=1): + for _ in range(int(STEP * 20)): + records = store.list_records()["items"] + if len(records) >= expected: + return records + await asyncio.sleep(0.05) + self.fail(f"审计没有落库：{store.list_records()}") + + async def hang_up(self, gate, endpoint, sent, hangup=None): + """发出聚合请求，等它卡在上游之后再让「客户端」挂断。""" + hangup = hangup or asyncio.Event() + task = asyncio.ensure_future( + gate(self.scope(endpoint), self.receive(self.hang_payload(endpoint), hangup), + self.collect(sent))) + await asyncio.wait_for(self.upstream_seen.wait(), STEP) + self.assertFalse(self.read_cancelled.is_set(), "客户端还没走，这枪不该结束") + hangup.set() + await asyncio.wait_for(task, STEP) # 修复前：没人监听断连，这一句必然超时 + return task + + async def succeed_after(self, gate, endpoint, sent): + """同一个名额闸上再打一发正常请求：名额必须已经还给网关。""" + await asyncio.wait_for( + gate(self.scope(endpoint), self.receive(self.payload(endpoint), asyncio.Event()), + self.collect(sent)), STEP) + + def store_for(self, name): + store = AuditStore(self.root / (name + ".sqlite3")) + self.addCleanup(store.close) + return store + + # --- 盲区本体：三个协议的聚合端点同一口径 --- + def test_hangup_during_aggregation_cancels_and_closes_the_upstream_call(self): + for endpoint in ENDPOINTS: + with self.subTest(endpoint=endpoint): + self.reset_upstream() + sent = [] + + async def scenario(): + await self.hang_up(self.gated(1), endpoint, sent) + + asyncio.run(scenario()) + self.assertTrue(self.read_cancelled.is_set(), "挂断之后，挂着的上游读要被取消") + self.assertTrue(self.stream_closed.is_set(), + "响应的异步 aclose() 也要跑完，连接才不留半开") + statuses = self.started_statuses(sent) + self.assertTrue(all(status in QUIET_STATUSES for status in statuses), statuses) + + def test_hangup_is_audited_as_cancelled_and_frees_the_slot(self): + for endpoint in ENDPOINTS: + with self.subTest(endpoint=endpoint): + self.reset_upstream() + store = self.store_for("hangup-" + endpoint) + + hung_up = [] + + async def scenario(): + gate = self.gated(1, store) + await self.hang_up(gate, endpoint, []) + hung_up.append((await self.drain_records(store))[0]) + await self.succeed_after(gate, endpoint, []) + + asyncio.run(scenario()) + record = hung_up[0] + self.assertEqual(record["outcome"], "cancelled", record) + self.assertFalse(record["streaming"], record) + + # --- 换凭证重放途中挂断：整段重放是一个可取消单元 --- + def test_hangup_after_credential_failover_cancels_the_whole_sequence(self): + for endpoint in ENDPOINTS: + with self.subTest(endpoint=endpoint): + self.reset_upstream() + self.fail_over_first = True + store = self.store_for("failover-" + endpoint) + sent, follow_up, hung_up = [], [], [] + + async def scenario(): + gate = self.gated(1, store) + with allow_failover(1): + await self.hang_up(gate, endpoint, sent) + hung_up.append((await self.drain_records(store))[0]) + await self.succeed_after(gate, endpoint, follow_up) + + asyncio.run(scenario()) + record = hung_up[0] + self.assertEqual(self.hang_attempts, 2, "第一枪 429 之后应恰好换凭证重放一次") + self.assertEqual(len(set(self.hang_uids)), 2, + "重放必须换凭证：" + str(self.hang_uids)) + self.assertTrue(self.read_cancelled.is_set()) + self.assertTrue(self.stream_closed.is_set(), "取消要等重放那一枪的响应也关干净") + self.assertTrue(all(status in QUIET_STATUSES + for status in self.started_statuses(sent)), sent) + self.assertEqual(record["outcome"], "cancelled", record) + self.assertNotIn("failover_recovered", json.dumps(record["attempts"] or []), + "没换回结果就不能标成已恢复") + self.assertEqual(self.started_statuses(follow_up), [200], + "挂断之后名额必须立刻可用") + + # --- 调用方自己取消外层任务：同样不许把上游留在半关状态 --- + def test_outer_task_cancellation_cancels_and_closes_the_upstream_call(self): + for endpoint in ENDPOINTS: + with self.subTest(endpoint=endpoint): + self.reset_upstream() + store = self.store_for("outer-" + endpoint) + sent, follow_up, hung_up = [], [], [] + + async def scenario(): + gate = self.gated(1, store) + task = asyncio.ensure_future( + gate(self.scope(endpoint), + self.receive(self.hang_payload(endpoint), asyncio.Event()), + self.collect(sent))) + await asyncio.wait_for(self.upstream_seen.wait(), STEP) + task.cancel() + with suppress(asyncio.CancelledError): + await task + hung_up.append((await self.drain_records(store))[0]) + await self.succeed_after(gate, endpoint, follow_up) + + asyncio.run(scenario()) + self.assertTrue(self.read_cancelled.is_set(), "外层取消要传到挂着的上游读") + self.assertTrue(self.stream_closed.is_set()) + self.assertEqual(hung_up[0]["outcome"], "cancelled", hung_up[0]) + self.assertEqual(self.started_statuses(follow_up), [200], follow_up) + + # --- 对照一：客户端没走，聚合请求必须照常完成（新监听不许误伤） --- + def test_aggregation_completes_while_the_client_is_still_there(self): + for endpoint in ENDPOINTS: + with self.subTest(endpoint=endpoint): + self.reset_upstream() + sent = [] + + async def scenario(): + await asyncio.wait_for( + self.gated(4)(self.scope(endpoint), + self.receive(self.payload(endpoint), asyncio.Event()), + self.collect(sent)), STEP) + + asyncio.run(scenario()) + self.assertEqual(self.started_statuses(sent), [200], sent) + self.assertIn(b"ok", b"".join(m.get("body", b"") for m in sent)) + self.assertFalse(self.read_cancelled.is_set(), "正常请求不该被断连监听打断") + + # --- 对照二：流式的同一场景在修复之前就已经成立 --- + def test_streaming_hangup_already_closes_the_upstream(self): + for endpoint in ENDPOINTS: + with self.subTest(endpoint=endpoint): + self.reset_upstream() + sent = [] + + async def scenario(): + gate = self.gated(1) + hangup = asyncio.Event() + task = asyncio.ensure_future( + gate(self.scope(endpoint), + self.receive(self.hang_payload(endpoint, stream=True), hangup), + self.collect(sent))) + await asyncio.wait_for(self.upstream_seen.wait(), STEP) + hangup.set() + with suppress(asyncio.CancelledError): + await asyncio.wait_for(task, STEP) + + asyncio.run(scenario()) + self.assertTrue(self.read_cancelled.is_set(), + "Starlette 的断连监听管的就是这一段") + self.assertIn(self.started_statuses(sent), ([], [200]), sent) + self.assertNotIn(b"[DONE]", b"".join(m.get("body", b"") for m in sent)) + + +if __name__ == "__main__": + unittest.main()