diff --git a/Sources/Functions/FunctionsClient.swift b/Sources/Functions/FunctionsClient.swift index 7bab76687..12521d1c8 100644 --- a/Sources/Functions/FunctionsClient.swift +++ b/Sources/Functions/FunctionsClient.swift @@ -272,6 +272,13 @@ public final class FunctionsClient: Sendable { _ functionName: String, options invokeOptions: FunctionInvokeOptions = .init() ) -> AsyncThrowingStream { + streamResponse(functionName, options: invokeOptions).stream + } + + func streamResponse( + _ functionName: String, + options invokeOptions: FunctionInvokeOptions + ) -> (stream: AsyncThrowingStream, delegate: StreamResponseDelegate) { let (stream, continuation) = AsyncThrowingStream.makeStream() let delegate = StreamResponseDelegate(continuation: continuation) @@ -284,13 +291,10 @@ public final class FunctionsClient: Sendable { task.resume() continuation.onTermination = { _ in - task.cancel() - - // Hold a strong reference to delegate until continuation terminates. - _ = delegate + session.invalidateAndCancel() } - return stream + return (stream, delegate) } private func buildRequest(functionName: String, options: FunctionInvokeOptions) diff --git a/Tests/FunctionsTests/FunctionsClientTests.swift b/Tests/FunctionsTests/FunctionsClientTests.swift index f0ac41fd6..d910c7e11 100644 --- a/Tests/FunctionsTests/FunctionsClientTests.swift +++ b/Tests/FunctionsTests/FunctionsClientTests.swift @@ -566,4 +566,32 @@ struct FunctionsClientTests { } catch FunctionsError.relayError { } } + + @Test + func invokeWithStreamedResponseInvalidatesSession() async throws { + let sut = makeSUT() + + Mock( + url: url.appendingPathComponent("stream"), + statusCode: 200, + data: [.post: Data("hello world".utf8)] + ) + .register() + + weak var weakDelegate: StreamResponseDelegate? + + do { + let (stream, delegate) = sut.streamResponse("stream", options: FunctionInvokeOptions()) + weakDelegate = delegate + for try await _ in stream {} + } + + // URLSession releases its delegate asynchronously after invalidation. + let deadline = Date().addingTimeInterval(5) + while weakDelegate != nil, Date() < deadline { + try await Task.sleep(nanoseconds: 10_000_000) + } + + #expect(weakDelegate == nil, "URLSession was not invalidated; its delegate leaked") + } }