mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
refactor(compression): 请求压缩策略改为跟随客户端行为,响应支持gzip压缩
- 移除全局 ENABLE_REQUEST_COMPRESSION 配置,改为根据客户端 Content-Encoding 决定是否对上游请求体进行 gzip 压缩 - 非流式响应根据客户端 Accept-Encoding 返回 gzip 压缩的 JSON - ApiRequestContext 记录客户端编码偏好并透传至 handler 链路 - 新增 http_compression 模块统一处理压缩相关判断逻辑 - 上游请求头丢弃列表新增 content-encoding 防止客户端值泄露 - ensure_json_body 支持解压 gzip 编码的请求体
This commit is contained in:
@@ -1,11 +1,17 @@
|
||||
"""测试 handler 基础工具函数"""
|
||||
|
||||
import gzip
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from src.api.handlers.base.utils import (
|
||||
build_json_response_for_client,
|
||||
build_sse_headers,
|
||||
extract_cache_creation_tokens,
|
||||
filter_proxy_response_headers,
|
||||
resolve_client_accept_encoding,
|
||||
resolve_client_content_encoding,
|
||||
)
|
||||
|
||||
|
||||
@@ -135,3 +141,51 @@ class TestFilterProxyResponseHeaders:
|
||||
|
||||
assert result["X-Request-Id"] == "abc"
|
||||
assert result["Anthropic-RateLimit-Requests-Remaining"] == "10"
|
||||
|
||||
|
||||
class TestResolveClientEncoding:
|
||||
def test_content_encoding_prefers_hint(self) -> None:
|
||||
headers = {"content-encoding": "gzip"}
|
||||
result = resolve_client_content_encoding(headers, hinted_content_encoding="br")
|
||||
assert result == "br"
|
||||
|
||||
def test_content_encoding_fallback_to_headers(self) -> None:
|
||||
headers = {"Content-Encoding": "gzip"}
|
||||
result = resolve_client_content_encoding(headers)
|
||||
assert result == "gzip"
|
||||
|
||||
def test_accept_encoding_prefers_hint(self) -> None:
|
||||
headers = {"accept-encoding": "gzip"}
|
||||
result = resolve_client_accept_encoding(headers, hinted_accept_encoding="br")
|
||||
assert result == "br"
|
||||
|
||||
def test_accept_encoding_fallback_to_headers(self) -> None:
|
||||
headers = {"Accept-Encoding": "gzip, deflate"}
|
||||
result = resolve_client_accept_encoding(headers)
|
||||
assert result == "gzip, deflate"
|
||||
|
||||
|
||||
class TestBuildJsonResponseForClient:
|
||||
def test_returns_gzip_response_when_client_accepts_gzip(self) -> None:
|
||||
response = build_json_response_for_client(
|
||||
status_code=200,
|
||||
content={"ok": True},
|
||||
headers={"content-type": "application/json"},
|
||||
client_accept_encoding="gzip, deflate",
|
||||
)
|
||||
|
||||
assert response.headers.get("content-encoding") == "gzip"
|
||||
assert "accept-encoding" in response.headers.get("vary", "").lower()
|
||||
decompressed = gzip.decompress(bytes(response.body))
|
||||
assert json.loads(decompressed.decode("utf-8")) == {"ok": True}
|
||||
|
||||
def test_returns_plain_json_when_gzip_not_accepted(self) -> None:
|
||||
response = build_json_response_for_client(
|
||||
status_code=200,
|
||||
content={"ok": True},
|
||||
headers={"content-type": "application/json"},
|
||||
client_accept_encoding="gzip;q=0, deflate",
|
||||
)
|
||||
|
||||
assert response.headers.get("content-encoding") is None
|
||||
assert json.loads(bytes(response.body).decode("utf-8")) == {"ok": True}
|
||||
|
||||
91
tests/unit/test_api_request_context.py
Normal file
91
tests/unit/test_api_request_context.py
Normal file
@@ -0,0 +1,91 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from starlette.requests import Request
|
||||
|
||||
from src.api.base.context import ApiRequestContext
|
||||
|
||||
|
||||
def _build_request(headers: dict[str, str] | None = None) -> Request:
|
||||
header_items = [
|
||||
(str(key).encode("latin-1"), str(value).encode("latin-1"))
|
||||
for key, value in (headers or {}).items()
|
||||
]
|
||||
scope = {
|
||||
"type": "http",
|
||||
"http_version": "1.1",
|
||||
"method": "POST",
|
||||
"scheme": "http",
|
||||
"path": "/v1/messages",
|
||||
"raw_path": b"/v1/messages",
|
||||
"query_string": b"",
|
||||
"headers": header_items,
|
||||
"client": ("127.0.0.1", 12345),
|
||||
"server": ("testserver", 80),
|
||||
}
|
||||
|
||||
async def receive() -> dict[str, object]:
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
request = Request(scope, receive)
|
||||
request.state.perf_metrics = {}
|
||||
return request
|
||||
|
||||
|
||||
def _build_context(raw_body: bytes, headers: dict[str, str] | None = None) -> ApiRequestContext:
|
||||
request = _build_request(headers=headers)
|
||||
return ApiRequestContext(
|
||||
request=request,
|
||||
db=None, # type: ignore[arg-type]
|
||||
user=None,
|
||||
api_key=None,
|
||||
request_id="req_test",
|
||||
start_time=0.0,
|
||||
client_ip="127.0.0.1",
|
||||
user_agent="pytest",
|
||||
original_headers=headers or {},
|
||||
query_params={},
|
||||
raw_body=raw_body,
|
||||
)
|
||||
|
||||
|
||||
class TestApiRequestContextEnsureJsonBody:
|
||||
def test_decompresses_gzip_body(self) -> None:
|
||||
payload = {"message": "hello", "count": 2}
|
||||
raw_body = gzip.compress(json.dumps(payload).encode("utf-8"))
|
||||
context = _build_context(raw_body, headers={"content-encoding": "gzip"})
|
||||
|
||||
result = context.ensure_json_body()
|
||||
|
||||
assert result == payload
|
||||
|
||||
def test_rejects_invalid_gzip_body(self) -> None:
|
||||
context = _build_context(b"not-gzip-body", headers={"content-encoding": "gzip"})
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
context.ensure_json_body()
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail == "gzip 请求体解压失败"
|
||||
|
||||
def test_build_records_client_encoding_preferences(self) -> None:
|
||||
request = _build_request(
|
||||
headers={
|
||||
"content-encoding": "gzip",
|
||||
"accept-encoding": "gzip, deflate",
|
||||
}
|
||||
)
|
||||
context = ApiRequestContext.build(
|
||||
request=request,
|
||||
db=None, # type: ignore[arg-type]
|
||||
user=None,
|
||||
api_key=None,
|
||||
raw_body=b"{}",
|
||||
)
|
||||
|
||||
assert context.client_content_encoding == "gzip"
|
||||
assert context.client_accept_encoding == "gzip, deflate"
|
||||
@@ -101,6 +101,7 @@ class TestBuildUpstreamHeaders:
|
||||
"X-Api-Key": "client",
|
||||
"User-Agent": "ua",
|
||||
"Content-Type": "text/plain",
|
||||
"Content-Encoding": "gzip",
|
||||
},
|
||||
"openai:chat",
|
||||
"provider",
|
||||
@@ -112,6 +113,7 @@ class TestBuildUpstreamHeaders:
|
||||
assert result["Authorization"] == "Bearer provider"
|
||||
assert result["User-Agent"] == "extra"
|
||||
assert result["Content-Type"] == "text/plain"
|
||||
assert "Content-Encoding" not in result
|
||||
assert result["X-Endpoint"] == "1"
|
||||
assert result["X-Extra"] == "1"
|
||||
|
||||
|
||||
48
tests/unit/test_proxy_resolver_compression.py
Normal file
48
tests/unit/test_proxy_resolver_compression.py
Normal file
@@ -0,0 +1,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
import json
|
||||
|
||||
from src.services.proxy_node.resolver import build_post_kwargs, build_stream_kwargs
|
||||
|
||||
|
||||
class TestProxyResolverCompression:
|
||||
def test_build_post_kwargs_compresses_when_client_sent_gzip(self) -> None:
|
||||
payload = {"message": "hello", "tokens": [1, 2, 3]}
|
||||
kwargs = build_post_kwargs(
|
||||
None,
|
||||
url="https://example.com/v1/messages",
|
||||
headers={"Content-Type": "application/json"},
|
||||
payload=payload,
|
||||
timeout=10.0,
|
||||
client_content_encoding="gzip",
|
||||
)
|
||||
|
||||
assert kwargs["headers"]["Content-Encoding"] == "gzip"
|
||||
assert json.loads(gzip.decompress(kwargs["content"]).decode("utf-8")) == payload
|
||||
|
||||
def test_build_post_kwargs_keeps_plain_body_without_client_gzip(self) -> None:
|
||||
payload = {"message": "plain"}
|
||||
kwargs = build_post_kwargs(
|
||||
None,
|
||||
url="https://example.com/v1/messages",
|
||||
headers={"Content-Type": "application/json"},
|
||||
payload=payload,
|
||||
timeout=10.0,
|
||||
client_content_encoding=None,
|
||||
)
|
||||
|
||||
assert "Content-Encoding" not in kwargs["headers"]
|
||||
assert json.loads(kwargs["content"].decode("utf-8")) == payload
|
||||
|
||||
def test_build_stream_kwargs_drops_stale_content_encoding_header(self) -> None:
|
||||
kwargs = build_stream_kwargs(
|
||||
None,
|
||||
url="https://example.com/v1/messages",
|
||||
headers={"content-encoding": "gzip", "Content-Type": "application/json"},
|
||||
payload={"message": "no-gzip"},
|
||||
timeout=10.0,
|
||||
client_content_encoding=None,
|
||||
)
|
||||
|
||||
assert all(key.lower() != "content-encoding" for key in kwargs["headers"])
|
||||
Reference in New Issue
Block a user