core/lmproto: await streaming AsyncResponse close

Support both sync and async close/aclose so stream cleanup does not
leave un-awaited coroutines.
This commit is contained in:
2026-07-14 23:24:05 +08:00
parent bd6d98ceda
commit a1cfc54f8e
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import inspect
from typing import TYPE_CHECKING, Literal, overload from typing import TYPE_CHECKING, Literal, overload
import msgspec import msgspec
@@ -24,6 +25,18 @@ if TYPE_CHECKING:
from niquests.models import AsyncResponse from niquests.models import AsyncResponse
async def _close_async_response(resp: AsyncResponse) -> None:
"""Close a streaming response; support sync or async close()."""
for name in ("aclose", "close"):
method = getattr(resp, name, None)
if not callable(method):
continue
result = method()
if inspect.isawaitable(result):
await result
return
class BaseOpenAIClient: class BaseOpenAIClient:
session: AsyncSession session: AsyncSession
encoder: msgspec.json.Encoder encoder: msgspec.json.Encoder
@@ -47,13 +60,7 @@ class BaseOpenAIClient:
if line.startswith(b"data: "): if line.startswith(b"data: "):
yield self.chunk_decoder.decode(line[6:]) yield self.chunk_decoder.decode(line[6:])
finally: finally:
aclose = getattr(resp, "aclose", None) await _close_async_response(resp)
if callable(aclose):
await aclose() # pyright: ignore[reportGeneralTypeIssues]
else:
close = getattr(resp, "close", None)
if callable(close):
_ = close()
class OpenAIClient(BaseOpenAIClient): class OpenAIClient(BaseOpenAIClient):