refactor: centralize LLM HTTP client lifecycle

Author: jiwenxu025-boop

scripts/bazi_engine/_http.py

@@ -1,15 +1,95 @@
-"""共享 HTTP 客户端 — 复用连接池，避免每个模块各自创建"""
-from contextlib import contextmanager
+"""共享 HTTP 客户端和调用级超时控制。"""
+import threading
+from collections.abc import Iterator
+from contextlib import asynccontextmanager, contextmanager
 
 import httpx
 
 _client: httpx.Client | None = None
+_async_client: httpx.AsyncClient | None = None
+_client_lock = threading.Lock()
 
 
-@contextmanager
-def shared_client(timeout: float = 60.0):
-    """获取共享的 httpx.Client（懒加载，单例，contextmanager 兼容 with 语句）"""
+class _TimeoutBoundClient:
+    """Apply a timeout to one call without mutating the shared connection pool."""
+
+    def __init__(self, client: httpx.Client, timeout: float):
+        self._client = client
+        self._timeout = timeout
+
+    def post(self, url: str, **kwargs) -> httpx.Response:
+        kwargs.setdefault("timeout", self._timeout)
+        return self._client.post(url, **kwargs)
+
+    def stream(self, method: str, url: str, **kwargs):
+        kwargs.setdefault("timeout", self._timeout)
+        return self._client.stream(method, url, **kwargs)
+
+
+class _AsyncTimeoutBoundClient:
+    """Async counterpart that shares one connection pool across chat streams."""
+
+    def __init__(self, client: httpx.AsyncClient, timeout: float):
+        self._client = client
+        self._timeout = timeout
+
+    async def post(self, url: str, **kwargs) -> httpx.Response:
+        kwargs.setdefault("timeout", self._timeout)
+        return await self._client.post(url, **kwargs)
+
+    def stream(self, method: str, url: str, **kwargs):
+        kwargs.setdefault("timeout", self._timeout)
+        return self._client.stream(method, url, **kwargs)
+
+
+def _get_client() -> httpx.Client:
     global _client
-    if _client is None:
-        _client = httpx.Client(timeout=timeout)
-    yield _client
+    with _client_lock:
+        if _client is None or _client.is_closed:
+            _client = httpx.Client(
+                limits=httpx.Limits(max_connections=12, max_keepalive_connections=6),
+            )
+        return _client
+
+
+def _get_async_client() -> httpx.AsyncClient:
+    global _async_client
+    if _async_client is None or _async_client.is_closed:
+        _async_client = httpx.AsyncClient(
+            limits=httpx.Limits(max_connections=12, max_keepalive_connections=6),
+        )
+    return _async_client
+
+
+def close_shared_client() -> None:
+    """Close the connection pool during application shutdown."""
+    global _client
+    with _client_lock:
+        client, _client = _client, None
+    if client is not None and not client.is_closed:
+        client.close()
+
+
+async def close_shared_clients() -> None:
+    """Close both pools during application shutdown."""
+    global _async_client
+    close_shared_client()
+    async_client, _async_client = _async_client, None
+    if async_client is not None and not async_client.is_closed:
+        await async_client.aclose()
+
+
+@contextmanager
+def shared_client(timeout: float = 60.0) -> Iterator[_TimeoutBoundClient]:
+    """Yield a shared pool with a timeout scoped to this individual request."""
+    if timeout <= 0:
+        raise ValueError("timeout must be positive")
+    yield _TimeoutBoundClient(_get_client(), timeout)
+
+
+@asynccontextmanager
+async def shared_async_client(timeout: float = 60.0):
+    """Yield the async connection pool with a timeout scoped to this request."""
+    if timeout <= 0:
+        raise ValueError("timeout must be positive")
+    yield _AsyncTimeoutBoundClient(_get_async_client(), timeout)

scripts/bazi_engine/api.py

@@ -22,7 +22,7 @@
 import threading
 import time
 import uuid
-from contextlib import suppress
+from contextlib import asynccontextmanager, suppress
 from datetime import datetime as dt
 from datetime import timedelta
 from pathlib import Path
@@ -34,6 +34,7 @@
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, Field
 
+from ._http import close_shared_clients
 from ._version import __version__
 from .chart import build_chart
 
@@ -77,7 +78,15 @@ class ActivationCodeRequest(BaseModel):
     code: str = Field(min_length=1, max_length=80)
 
 
-app = FastAPI(title="八字排盘引擎", version=__version__)
+@asynccontextmanager
+async def lifespan(_app: FastAPI):
+    try:
+        yield
+    finally:
+        await close_shared_clients()
+
+
+app = FastAPI(title="八字排盘引擎", version=__version__, lifespan=lifespan)
 
 # 前端页面路径（支持环境变量覆盖，打包部署时可指定）
 _frontend_env = os.getenv("FRONTEND_DIR", "")

scripts/bazi_engine/chat.py

@@ -12,10 +12,12 @@
 
 import httpx
 
+from ._deepseek_config import DEEPSEEK_API_URL, DEEPSEEK_KEY, DEEPSEEK_MODEL
+from ._http import shared_async_client
+
 # ═══════════════════════════════════════════════════════════════
 # 配置
 # ═══════════════════════════════════════════════════════════════
-from ._deepseek_config import DEEPSEEK_API_URL, DEEPSEEK_KEY, DEEPSEEK_MODEL
 
 _ACTIVATION_FILE = Path(__file__).resolve().parent / "activation_codes.json"
 _RUNTIME_DATA_LOCK = threading.RLock()
@@ -334,7 +336,7 @@ async def call_deepseek_stream(messages: list[dict]) -> AsyncGenerator[str]:
 
     try:
         async with (
-            httpx.AsyncClient(timeout=60.0) as client,
+            shared_async_client(60.0) as client,
             client.stream("POST", DEEPSEEK_API_URL, json=payload, headers=headers) as resp,
         ):
             if resp.status_code != 200:

scripts/tests/test_http.py

@@ -0,0 +1,77 @@
+"""Shared HTTP client lifecycle and timeout tests."""
+
+import asyncio
+
+import httpx
+import pytest
+from fastapi.testclient import TestClient
+
+
+def test_shared_client_keeps_timeouts_scoped_to_each_request(monkeypatch):
+    import bazi_engine._http as http_module
+
+    seen_timeouts = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen_timeouts.append(request.extensions["timeout"])
+        return httpx.Response(200, json={"ok": True})
+
+    raw_client = httpx.Client(transport=httpx.MockTransport(handler))
+    monkeypatch.setattr(http_module, "_client", raw_client)
+
+    with http_module.shared_client(5.0) as client:
+        assert client.post("https://example.test/first").status_code == 200
+    with http_module.shared_client(12.0) as client:
+        assert client.post("https://example.test/second").status_code == 200
+
+    assert seen_timeouts[0]["read"] == 5.0
+    assert seen_timeouts[1]["read"] == 12.0
+    http_module.close_shared_client()
+
+
+def test_shared_client_rejects_non_positive_timeout():
+    import bazi_engine._http as http_module
+
+    with pytest.raises(ValueError, match="timeout must be positive"), http_module.shared_client(0):
+        pass
+
+
+def test_api_lifespan_closes_shared_http_client(monkeypatch):
+    import bazi_engine._http as http_module
+    from bazi_engine.api import app
+
+    raw_client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200)))
+    raw_async_client = httpx.AsyncClient(transport=httpx.MockTransport(lambda _request: httpx.Response(200)))
+    monkeypatch.setattr(http_module, "_client", raw_client)
+    monkeypatch.setattr(http_module, "_async_client", raw_async_client)
+
+    with TestClient(app) as client:
+        assert client.get("/api/health").status_code == 200
+
+    assert raw_client.is_closed
+    assert raw_async_client.is_closed
+
+
+def test_shared_async_client_keeps_timeouts_scoped_to_each_request(monkeypatch):
+    import bazi_engine._http as http_module
+
+    seen_timeouts = []
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        seen_timeouts.append(request.extensions["timeout"])
+        return httpx.Response(200, json={"ok": True})
+
+    raw_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+    monkeypatch.setattr(http_module, "_async_client", raw_client)
+
+    async def make_calls():
+        async with http_module.shared_async_client(7.0) as client:
+            assert (await client.post("https://example.test/first")).status_code == 200
+        async with http_module.shared_async_client(13.0) as client:
+            assert (await client.post("https://example.test/second")).status_code == 200
+
+    asyncio.run(make_calls())
+
+    assert seen_timeouts[0]["read"] == 7.0
+    assert seen_timeouts[1]["read"] == 13.0
+    asyncio.run(http_module.close_shared_clients())