Python: [Bug]: an output guard on a streamed run releases the full answer before it can raise, unless stream_buffer_updates=True is set #8954

# agent-framework: an output guard in agent middleware does not stop a streamed answer
# from reaching the consumer unless the stream is explicitly buffered.
# Offline: scripted chat client, nothing leaves the process.
import asyncio
import warnings
from importlib.metadata import version

from agent_framework import (Agent, AgentMiddleware, BaseChatClient, ChatMiddlewareLayer, ChatResponse,
                             ChatResponseUpdate, Content, FunctionInvocationLayer, Message, ResponseStream)
from agent_framework.observability import ChatTelemetryLayer

PIECES = ["the secret ", "is ", "1234"]
EXPECTED = "'' / guard raised"          # the reader should receive nothing


class Answers(FunctionInvocationLayer, ChatMiddlewareLayer, ChatTelemetryLayer, BaseChatClient):
    """Model that answers "the secret is 1234" (in three pieces when streamed)."""

    def __init__(self):
        super().__init__(middleware=[])

    def _inner_get_response(self, *, messages, stream, options, **kw):
        if stream:
            async def updates():
                for piece in PIECES:
                    yield ChatResponseUpdate(contents=[Content.from_text(piece)], role="assistant")
            return ResponseStream(updates(), finalizer=ChatResponse.from_updates)

        async def whole():
            return ChatResponse(messages=Message(role="assistant", contents=["".join(PIECES)]))
        return whole()


def check(final):
    """The output guard: refuse any answer that contains the secret."""
    if "1234" in final.text:
        raise RuntimeError("blocked by output guard")
    return final


class ResultHookGuard(AgentMiddleware):
    """Checks the finished answer with with_result_hook (works on 1.19.0 and main)."""

    async def process(self, context, call_next):
        await call_next()
        if isinstance(context.result, ResponseStream):
            context.result = context.result.with_result_hook(check)
        else:
            check(context.result)


class ResultGateGuard(AgentMiddleware):
    """Checks the finished answer with a result gate, the API documented on main for blocking output."""

    def __init__(self, buffer):
        self.buffer = buffer

    async def process(self, context, call_next):
        if context.stream:
            if not hasattr(context, "stream_result_gates_after"):
                raise NotImplementedError
            context.stream_result_gates_after.append(check)      # registered before call_next()
            context.stream_buffer_updates = self.buffer
            await call_next()
        else:
            await call_next()
            check(context.result)


async def run(guard, stream):
    delivered = []
    try:
        if stream:
            async for update in Agent(client=Answers(), middleware=[guard]).run("go", stream=True):
                if update.text:
                    delivered.append(update.text)       # what the application forwards to the reader
        else:
            delivered.append((await Agent(client=Answers(), middleware=[guard]).run("go")).text)
    except NotImplementedError:
        return "n/a (API not in this release)"
    except RuntimeError:
        return f"{''.join(delivered)!r:20} / guard raised"
    return f"{''.join(delivered)!r:20} / guard did NOT raise"


async def main():
    print("agent-framework-core", version("agent-framework-core"))
    print(f"{'scenario':48} {'expected':24} actual")
    rows = [("not streamed, result hook", ResultHookGuard(), False),
            ("streamed, with_result_hook", ResultHookGuard(), True),
            ("streamed, result gate", ResultGateGuard(buffer=False), True),
            ("streamed, result gate + stream_buffer_updates", ResultGateGuard(buffer=True), True)]
    for name, guard, stream in rows:
        print(f"{name:48} {EXPECTED:24} {await run(guard, stream)}")


warnings.simplefilter("ignore")
asyncio.run(main())