Python: close native generators when stream invocations end

python/semantic_kernel/functions/kernel_function.py

@@ -5,6 +5,7 @@
 import time
 from abc import abstractmethod
 from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
+from contextlib import aclosing, closing
 from copy import copy, deepcopy
 from inspect import isasyncgen, isgenerator
 from typing import TYPE_CHECKING, Any
@@ -352,13 +353,15 @@ async def invoke_stream(
                 function_results: list[Any] = []
                 if function_context.result is not None:
                     if isasyncgen(function_context.result.value):
-                        async for partial in function_context.result.value:
-                            function_results.append(partial)
-                            yield partial
+                        async with aclosing(function_context.result.value) as async_stream:
+                            async for partial in async_stream:
+                                function_results.append(partial)
+                                yield partial
                     elif isgenerator(function_context.result.value):
-                        for partial in function_context.result.value:
-                            function_results.append(partial)
-                            yield partial
+                        with closing(function_context.result.value) as stream:
+                            for partial in stream:
+                                function_results.append(partial)
+                                yield partial
                     else:
                         function_results.append(function_context.result.value)
                         yield function_context.result

python/semantic_kernel/kernel.py

@@ -2,7 +2,7 @@
 
 import logging
 from collections.abc import AsyncGenerator, AsyncIterable, Callable
-from contextlib import AbstractAsyncContextManager
+from contextlib import AbstractAsyncContextManager, aclosing
 from copy import copy, deepcopy
 from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
@@ -142,15 +142,16 @@ async def invoke_stream(
 
         function_result: list[list["StreamingContentMixin"] | Any] = []
 
-        async for stream_message in function.invoke_stream(self, arguments):
-            if isinstance(stream_message, FunctionResult) and (
-                exception := stream_message.metadata.get(METADATA_EXCEPTION_KEY, None)
-            ):
-                raise KernelInvokeException(
-                    f"Error occurred while invoking function: '{function.fully_qualified_name}'"
-                ) from exception
-            function_result.append(stream_message)
-            yield stream_message
+        async with aclosing(function.invoke_stream(self, arguments)) as stream:
+            async for stream_message in stream:
+                if isinstance(stream_message, FunctionResult) and (
+                    exception := stream_message.metadata.get(METADATA_EXCEPTION_KEY, None)
+                ):
+                    raise KernelInvokeException(
+                        f"Error occurred while invoking function: '{function.fully_qualified_name}'"
+                    ) from exception
+                function_result.append(stream_message)
+                yield stream_message
 
         if return_function_results:
             output_function_result: list["StreamingContentMixin"] = []

python/tests/unit/functions/test_kernel_function_from_method.py

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 import logging
+from asyncio import CancelledError
 from collections.abc import AsyncGenerator, Iterable
 from contextlib import nullcontext
 from dataclasses import dataclass
@@ -26,6 +27,85 @@
 kernel_function_module = import_module("semantic_kernel.functions.kernel_function")
 
 
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("via_kernel", [False, True])
+@pytest.mark.parametrize("termination", ["close", "cancel", "exhaust", "error"])
+async def test_invoke_stream_closes_native_generator(
+    kernel: Kernel, asynchronous: bool, via_kernel: bool, termination: str
+):
+    closed = []
+    failure = ValueError("native stream failed")
+
+    @kernel_function
+    def sync_stream() -> Iterable[str]:
+        try:
+            yield "first"
+            if termination == "error":
+                raise failure
+            yield "second"
+        finally:
+            closed.append(True)
+
+    @kernel_function
+    async def async_stream() -> AsyncGenerator[str, None]:
+        try:
+            yield "first"
+            if termination == "error":
+                raise failure
+            yield "second"
+        finally:
+            closed.append(True)
+
+    # Retain the source so garbage collection cannot stand in for explicit cleanup.
+    method = async_stream if asynchronous else sync_stream
+    source = method()
+    function = KernelFunctionFromMethod(method=method, stream_method=lambda: source)
+    stream = kernel.invoke_stream(function) if via_kernel else function.invoke_stream(kernel)
+    try:
+        assert await anext(stream) == "first"
+        assert not closed
+        if termination == "close":
+            await stream.aclose()
+        elif termination == "cancel":
+            with pytest.raises(CancelledError):
+                await stream.athrow(CancelledError())
+        elif termination == "error":
+            with pytest.raises(ValueError) as exc_info:
+                await anext(stream)
+            assert exc_info.value is failure
+        else:
+            assert [part async for part in stream] == ["second"]
+        assert closed == [True]
+    finally:
+        await stream.aclose()
+        if asynchronous:
+            await source.aclose()
+        else:
+            source.close()
+
+
+@pytest.mark.parametrize("via_kernel", [False, True])
+async def test_invoke_stream_does_not_close_ordinary_result(kernel: Kernel, via_kernel: bool):
+    closed = []
+
+    class ResultList(list):
+        def close(self):
+            closed.append(True)
+
+    @kernel_function
+    def native() -> str:
+        return "non-streaming result"
+
+    result = ResultList(["ordinary result"])
+    function = KernelFunctionFromMethod(method=native, stream_method=lambda: result)
+    stream = kernel.invoke_stream(function) if via_kernel else function.invoke_stream(kernel)
+    parts = [part async for part in stream]
+    assert len(parts) == 1
+    assert parts[0].value is result
+    assert result == ["ordinary result"]
+    assert not closed
+
+
 class CustomType(KernelBaseModel):
     id: str
     name: str