Streaming responses never return their connection to the pool (closed before the chunked terminator is read) #4040

# AsyncStream.__stream__
try:
    async for sse in iterator:
        if sse.data.startswith("[DONE]"):
            break
        ...
finally:
    await response.aclose()

if sse.data.startswith("[DONE]"):
    # The server has signalled completion; only the body terminator remains.
    # Read it so the connection can return to the pool.
    async for _ in iterator:
        pass
    break

import asyncio
import json
import re
import sys
import threading

import openai
from openai import AsyncOpenAI, OpenAI

CHUNKS = [
    {"id": "c", "object": "chat.completion.chunk", "created": 0, "model": "m",
     "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]},
    {"id": "c", "object": "chat.completion.chunk", "created": 0, "model": "m",
     "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
]
COMPLETION = {"id": "c", "object": "chat.completion", "created": 0, "model": "m",
              "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"},
                           "finish_reason": "stop"}]}
MSGS = [{"role": "user", "content": "hi"}]
connections = 0


async def serve(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    global connections
    connections += 1
    try:
        while True:
            head = await reader.readuntil(b"\r\n\r\n")
            length = int(re.search(rb"content-length: (\d+)", head, re.I).group(1))
            body = json.loads(await reader.readexactly(length))
            if body.get("stream"):
                writer.write(b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n"
                             b"transfer-encoding: chunked\r\n\r\n")
                events = [f"data: {json.dumps(c)}\n\n".encode() for c in CHUNKS]
                events.append(b"data: [DONE]\n\n")
                for ev in events:
                    writer.write(f"{len(ev):x}\r\n".encode() + ev + b"\r\n")
                    await writer.drain()
                    await asyncio.sleep(0.005)
                writer.write(b"0\r\n\r\n")
            else:
                payload = json.dumps(COMPLETION).encode()
                writer.write(b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\n"
                             + f"content-length: {len(payload)}\r\n\r\n".encode() + payload)
            await writer.drain()
    except (asyncio.IncompleteReadError, ConnectionError):
        pass
    finally:
        writer.close()


def start_server() -> tuple[int, callable]:
    ready = threading.Event()
    state: dict = {}

    def run() -> None:
        async def main() -> None:
            server = await asyncio.start_server(serve, "127.0.0.1", 0)
            state["port"] = server.sockets[0].getsockname()[1]
            state["loop"] = asyncio.get_running_loop()
            state["stop"] = asyncio.Event()
            ready.set()
            await state["stop"].wait()
            server.close()
            server.close_clients()
            await server.wait_closed()

        asyncio.run(main())

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    ready.wait()

    def stop() -> None:
        state["loop"].call_soon_threadsafe(state["stop"].set)
        thread.join(timeout=5)

    return state["port"], stop


def report(label: str) -> None:
    global connections
    print(f"{label:<44} -> {connections} TCP connection(s)")
    connections = 0


async def async_cases(base_url: str) -> None:
    client = AsyncOpenAI(api_key="x", base_url=base_url)
    for _ in range(5):
        async for _ in await client.chat.completions.create(model="m", messages=MSGS, stream=True):
            pass
    report("AsyncOpenAI streaming x5, fully consumed")

    client = AsyncOpenAI(api_key="x", base_url=base_url)
    for _ in range(5):
        async with await client.chat.completions.create(model="m", messages=MSGS, stream=True) as s:
            async for _ in s:
                pass
    report("AsyncOpenAI streaming x5, async with")

    client = AsyncOpenAI(api_key="x", base_url=base_url)
    for _ in range(5):
        await client.chat.completions.create(model="m", messages=MSGS)
    report("AsyncOpenAI non-streaming x5")

    try:
        import httpx
    except ImportError:
        return
    async with httpx.AsyncClient() as raw:
        for _ in range(5):
            async with raw.stream("POST", f"{base_url}/chat/completions",
                                  json={"model": "m", "messages": MSGS, "stream": True}) as r:
                async for _ in r.aiter_lines():
                    pass
    report(f"raw httpx {httpx.__version__} stream x5, read to EOF")


def sync_cases(base_url: str) -> None:
    client = OpenAI(api_key="x", base_url=base_url)
    for _ in range(5):
        for _ in client.chat.completions.create(model="m", messages=MSGS, stream=True):
            pass
    report("OpenAI (sync) streaming x5, fully consumed")

    client = OpenAI(api_key="x", base_url=base_url)
    for _ in range(5):
        client.chat.completions.create(model="m", messages=MSGS)
    report("OpenAI (sync) non-streaming x5")


if __name__ == "__main__":
    port, stop = start_server()
    base_url = f"http://127.0.0.1:{port}/v1"
    print(f"openai {openai.__version__}, python {sys.version.split()[0]}")
    asyncio.run(async_cases(base_url))
    sync_cases(base_url)
    stop()