Python: [Bug]: Ollama and Gemini chat clients leave the SDK stream open when the consumer stops early #9143

import asyncio
from unittest.mock import AsyncMock, MagicMock

from agent_framework import Message
from agent_framework_gemini import GeminiChatClient
from agent_framework_ollama import OllamaChatClient
from google.genai import types
from ollama._types import ChatResponse
from ollama._types import Message as OllamaMessage

closed = {}


def sdk_stream(name, chunks):
    # Stand-in for the SDK stream, which keeps the HTTP response open until it is closed.
    async def gen():
        try:
            for chunk in chunks:
                yield chunk
        finally:
            closed[name] = True

    closed[name] = False
    return gen()


async def stop_early(name, client):
    stream = client._inner_get_response(messages=[Message(role="user", contents=["hi"])], options={}, stream=True)
    async with stream:
        async for _ in stream:
            break  # the consumer stops after the first update
    print(f"{name}: SDK stream closed after the consumer stopped = {closed[name]}")


async def main():
    ollama_sdk = MagicMock()
    ollama_sdk.chat = AsyncMock(
        return_value=sdk_stream("ollama", [ChatResponse(message=OllamaMessage(role="assistant", content="hello"), model="m")])
    )
    await stop_early("ollama", OllamaChatClient(client=ollama_sdk, model="m"))

    genai = MagicMock()
    genai._api_client.vertexai = False
    genai._api_client._http_options.base_url = "https://generativelanguage.googleapis.com/"
    chunk = types.GenerateContentResponse(
        candidates=[types.Candidate(content=types.Content(role="model", parts=[types.Part(text="hello")]))]
    )
    genai.aio.models.generate_content_stream = AsyncMock(return_value=sdk_stream("gemini", [chunk]))
    await stop_early("gemini", GeminiChatClient(client=genai, model="gemini-2.5-flash"))


asyncio.run(main())

Actual:
ollama: SDK stream closed after the consumer stopped = False
gemini: SDK stream closed after the consumer stopped = False

Expected:
ollama: SDK stream closed after the consumer stopped = True
gemini: SDK stream closed after the consumer stopped = True