fix: close stream generators and cancel embeddings on disconnect

modelship/infer/base_infer.py

@@ -122,10 +122,23 @@ async def run_cancellable_stream(
         the in-flight `__anext__()` call delivers `CancelledError` straight into
         `work`'s currently-suspended frame (and transitively into whatever it's
         awaiting, e.g. an engine's own generator), the same way cancelling a
-        plain task does for `run_cancellable`. `aclose()` afterward is a defensive
-        no-op when the generator already self-terminated on that exception.
+        plain task does for `run_cancellable`. The `finally` block's `aclose()`
+        is what actually guarantees `work` is closed on every exit path —
+        including the consumer closing *this* generator early (`GeneratorExit`
+        propagating out of the `yield`), which the disconnect branch's own
+        `aclose()` above doesn't cover. It's a defensive no-op wherever `work`
+        already self-terminated.
+
+        `next_item` is tracked outside the loop so `finally` can reach it: if
+        this generator is torn down (cancelled or `aclose()`d) while suspended
+        in the `asyncio.wait` below rather than at the `yield`, `next_item`'s
+        `__anext__()` call is still in flight and still owns `work`'s frame —
+        calling `work.aclose()` before that settles raises `RuntimeError:
+        aclose(): asynchronous generator is already running`. It must be
+        cancelled and awaited first, same as the disconnect branch above does.
         """
         watch = asyncio.ensure_future(self._poll_disconnect(raw_request))
+        next_item: asyncio.Task[T] | None = None
         try:
             while True:
                 next_item = asyncio.ensure_future(work.__anext__())
@@ -147,6 +160,11 @@ async def run_cancellable_stream(
             watch.cancel()
             with contextlib.suppress(asyncio.CancelledError):
                 await watch
+            if next_item is not None and not next_item.done():
+                next_item.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await next_item
+            await work.aclose()
 
     async def on_generation_aborted(self) -> None:
         """Hook for loaders whose engine needs cleanup beyond task cancellation

modelship/infer/llama_server/llama_server_infer.py

@@ -22,6 +22,7 @@
     UnsupportedContentError,
     build_from_parsed,
     build_responses_items_from_parsed,
+    encode_chat_sse_chunk,
     normalize_chat_messages,
 )
 from modelship.openai.protocol import (
@@ -381,6 +382,14 @@ async def create_embedding(
         payload = request.model_dump(exclude_none=True, exclude={"model"})
         payload["model"] = self.model_config.name
 
+        try:
+            return await self.run_cancellable(self._send_embedding(payload, request_id), raw_request)
+        except ClientDisconnectedError:
+            logger.info("embedding request %s aborted: client disconnected", request_id)
+            return create_error_response("Client disconnected")
+
+    async def _send_embedding(self, payload: dict[str, Any], request_id: str) -> ErrorResponse | EmbeddingResponse:
+        assert self._client is not None
         try:
             resp = await self._client.post("/v1/embeddings", json=payload)
             resp.raise_for_status()
@@ -536,15 +545,17 @@ async def _stream_chat_completion(
 
     async def _stream_chat_completion_body(self, payload: dict[str, Any], request_id: str) -> AsyncGenerator[str, None]:
         assert self._client is not None
+        trace = logger.isEnabledFor(TRACE)
         buffered: list[str] = []
         try:
             async for chunk in _raw_stream_chunks(
                 self._client, payload, model_name=self.model_config.name, request_id=request_id
             ):
-                for choice in chunk.choices:
-                    if choice.delta.content:
-                        buffered.append(choice.delta.content)
-                yield _encode_chunk(chunk)
+                if trace:
+                    for choice in chunk.choices:
+                        if choice.delta.content:
+                            buffered.append(choice.delta.content)
+                yield encode_chat_sse_chunk(chunk)
             yield "data: [DONE]\n\n"
         except _LlamaServerStreamError as e:
             yield _encode_error(str(e))
@@ -554,7 +565,8 @@ async def _stream_chat_completion_body(self, payload: dict[str, Any], request_id
             yield _encode_error(f"llama-server request failed: {e}")
             yield "data: [DONE]\n\n"
         finally:
-            logger.log(TRACE, "chat response %s (stream): %r", request_id, "".join(buffered))
+            if trace:
+                logger.log(TRACE, "chat response %s (stream): %r", request_id, "".join(buffered))
 
     async def _stream_response(
         self, payload: dict[str, Any], request: ResponsesRequest, request_id: str, raw_request: RawRequestProxy
@@ -842,12 +854,7 @@ def _project_embedding_response(data: dict, *, model_name: str) -> EmbeddingResp
                 index=item.get("index", 0), embedding=item.get("embedding", []), object=item.get("object", "embedding")
             )
         )
-    raw_usage = data.get("usage") or {}
-    usage = UsageInfo(
-        prompt_tokens=raw_usage.get("prompt_tokens", 0) or 0,
-        completion_tokens=raw_usage.get("completion_tokens", 0) or 0,
-        total_tokens=raw_usage.get("total_tokens", 0) or 0,
-    )
+    usage = _project_usage(data.get("usage"))
     created = data.get("created")
     if created is None:
         created = int(time.time())
@@ -861,9 +868,5 @@ def _project_embedding_response(data: dict, *, model_name: str) -> EmbeddingResp
     )
 
 
-def _encode_chunk(chunk: ChatCompletionStreamResponse) -> str:
-    return f"data: {json.dumps(chunk.model_dump(mode='json'))}\n\n"
-
-
 def _encode_error(detail: str) -> str:
     return f"data: {json.dumps({'error': {'message': detail, 'type': 'api_error'}})}\n\n"

modelship/infer/vllm/vllm_infer.py

@@ -68,12 +68,12 @@
     UnsupportedContentError,
     build_from_parsed,
     build_responses_items_from_parsed,
+    encode_chat_sse_chunk,
     normalize_chat_messages,
 )
 from modelship.openai.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponse,
-    ChatCompletionStreamResponse,
     EmbeddingCompletionRequest,
     EmbeddingRequest,
     ErrorResponse,
@@ -102,10 +102,6 @@
 logger = get_logger("infer.vllm")
 
 
-def _encode_chunk(chunk: ChatCompletionStreamResponse) -> str:
-    return f"data: {json.dumps(chunk.model_dump(mode='json'))}\n\n"
-
-
 def _encode_error(error: ErrorResponse) -> str:
     return f"data: {json.dumps(error.model_dump(mode='json'))}\n\n"
 
@@ -441,13 +437,15 @@ async def _create_chat_completion_stream(
             want_logprobs=bool(request.logprobs),
             num_output_top_logprobs=request.top_logprobs,
         )
+        trace = logger.isEnabledFor(TRACE)
         buffered: list[str] = []
         try:
             async for chunk in self.run_cancellable_stream(stream, raw_request):
-                for choice in chunk.choices:
-                    if choice.delta.content:
-                        buffered.append(choice.delta.content)
-                yield _encode_chunk(chunk)
+                if trace:
+                    for choice in chunk.choices:
+                        if choice.delta.content:
+                            buffered.append(choice.delta.content)
+                yield encode_chat_sse_chunk(chunk)
         except ClientDisconnectedError:
             logger.info("chat request %s aborted: client disconnected", request_id)
             return
@@ -463,7 +461,8 @@ async def _create_chat_completion_stream(
             yield "data: [DONE]\n\n"
             return
         finally:
-            logger.log(TRACE, "chat response %s (stream): %r", request_id, "".join(buffered))
+            if trace:
+                logger.log(TRACE, "chat response %s (stream): %r", request_id, "".join(buffered))
         yield "data: [DONE]\n\n"
 
     async def _create_chat_completion_no_stream(

modelship/openai/chat_utils.py

@@ -8,6 +8,7 @@
 from modelship.openai.protocol import (
     ChatCompletionResponse,
     ChatCompletionResponseChoice,
+    ChatCompletionStreamResponse,
     ChatMessage,
     ToolCall,
     UsageInfo,
@@ -279,6 +280,11 @@ def build_from_parsed(
     )
 
 
+def encode_chat_sse_chunk(chunk: ChatCompletionStreamResponse) -> str:
+    """Encode one chat-completion stream chunk as an SSE `data:` line."""
+    return f"data: {chunk.model_dump_json()}\n\n"
+
+
 def build_responses_items_from_parsed(parsed: ParsedChatOutput) -> list[ResponseOutputItem]:
     """Shape one parsed choice into Responses ``output[]`` items.
 

tests/test_base_infer_disconnect.py

@@ -199,3 +199,67 @@ async def gen():
 
     assert items == [0, 1, 2, 3, 4]
     assert elapsed < _DISCONNECT_POLL_INTERVAL_S
+
+
+@pytest.mark.asyncio
+async def test_consumer_closing_outer_generator_early_closes_work():
+    """If the consumer stops iterating (e.g. Starlette aclose()s the response
+    generator mid-stream) before any disconnect is detected, `work` must still
+    be closed — not left for the event loop to finalize whenever it gets
+    around to it."""
+    work_closed = asyncio.Event()
+
+    async def gen():
+        try:
+            for i in range(5):
+                yield i
+                await asyncio.sleep(10)
+        finally:
+            work_closed.set()
+
+    infer = _Infer()
+    raw_request = _FakeRawRequest(disconnect_after=None)
+
+    outer = infer.run_cancellable_stream(gen(), raw_request)
+    first = await outer.__anext__()
+    assert first == 0
+
+    await outer.aclose()
+
+    assert work_closed.is_set()
+    assert infer.aborted is False
+
+
+@pytest.mark.asyncio
+async def test_task_cancelled_while_suspended_in_asyncio_wait_does_not_raise():
+    """If the task driving this generator is cancelled while suspended in the
+    internal `asyncio.wait` (rather than at the `yield`), the loop's `next_item`
+    task is still in flight and still owns `work`'s frame. Calling
+    `work.aclose()` before cancelling and awaiting `next_item` first raises
+    `RuntimeError: aclose(): asynchronous generator is already running`
+    instead of letting `CancelledError` propagate cleanly."""
+    work_closed = asyncio.Event()
+
+    async def gen():
+        try:
+            yield "first"
+            await asyncio.sleep(10)
+            yield "unreachable"  # pragma: no cover - never reached
+        finally:
+            work_closed.set()
+
+    infer = _Infer()
+    raw_request = _FakeRawRequest(disconnect_after=None)
+
+    outer = infer.run_cancellable_stream(gen(), raw_request)
+    first = await outer.__anext__()
+    assert first == "first"
+
+    task = asyncio.ensure_future(outer.__anext__())
+    await asyncio.sleep(0.01)  # let it create next_item and suspend in asyncio.wait
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert work_closed.is_set()