Stop Streaming Early and Save Partial Response to History

[CristinaStn]

I need to stop agent streaming early (e.g., user clicks "Stop" button) and save the partial response to the history provider. The use case is:

1. User sends a query
2. Agent starts streaming response
3. User cancels after receiving a few chunks
4. I want to save the truncated conversation (user message + partial assistant response) to Cosmos DB history

The Problem I Encountered

When I break the stream loop and call get_final_response(), it consumes the entire remaining stream:

foundry_client = FoundryChatClient(...)
agent = Agent(
                client=foundry_client,
                name="HistoryAgent",
                instructions="You are a helpful assistant",
                tools=[code_interpreter],
                context_providers=[history_provider],
                default_options={"store": True},
            )
session = agent.create_session()
response = agent.run(query, session=session, stream=True)

async for update in response:                                                                                                                                                             
    if user_clicked_stop:
        break  # Exit early                                                                                                                                                               

  # This consumes the rest of the stream and gives full response                                                                                                                       
final = await response.get_final_response()

Workaround I Found

Setting the internal ResponseStream._consumed flag before calling get_final_response():

async for update in response:                                                                                                                                                             
    collected_text.append(update.text)

    if len(collected_text) >= 2:  # Stop after 2 chunks                                                                                                                                   
        response._consumed = True  # Mark as consumed
        break    
# Now this uses only the collected updates
final = await response.get_final_response()                                                 

This works - both user and truncated assistant messages are saved to history correctly.
Concerns with my workaround:

• Uses private API (_consumed)
• Could break in future versions
• Not sure if all cleanup hooks run properly
• No official documentation to reference

Questions: is there:

1. An official API I missed for stopping streams?
2. A different pattern I should use?
3. Any reason this workaround might cause issues?

Environment

• Python 3.11
• agent-framework 1.1.2
• Azure Foundry Chat Client
• CosmosDB History Provider

[moonbox3]

Adding @eavanvalkenburg to comment

[eavanvalkenburg]

@CristinaStn thanks for posting, we could add a cancel method to the ResponseStream object, which would set _completed to True and probably does some other fixes, however the concern I have is that this runs a big risk of ending with the history in a bad state, which will cause new issues that are really hard to pinpoint, for instance, because if you happen to cancel the stream right after a tool call then it might be ok, while if you do it while the model is sending tool calls, you first get malformed json in the tool call content and then a model complaining that there isn't a valid pair of tool call and result, and the number of race conditions that this unlocks is vast, now we also have reasoning content attached to tool calls, which are another vector for this, etc.

The other problem is that the finalizers are designed to be flexible, so in some cases it just combines update into a single response, but it can also be used to update UI, to store in file storage, etc. which is all circumvented when setting _consumed to True without actually consuming.

So, the best alternative, is to use a FunctionMiddleware to interrupt before the model goes into the next round of model calls, instead returning, if you do the raise MiddelwareTermination after await call_next() then the history will stay in a good shape and you can fully consume the stream without gaining additional calls to the model. For other instances, just not handling the results but letting the stream complete is the way to go, a ChatMiddleware could also be used to set tool_mode to None so that you know it will not do another tool call, but forces the model to respond.

@westey-m has also done a lot of looking into different approaches here, he might have some other insights to share, but the long and short of it, is yes this will cause other issues that are much more tricky to troubleshoot/recover from!

[CristinaStn]

Thanks @eavanvalkenburg, sounds like FunctionMiddleware approach will be a better way to go. However it does not stop the streaming immediately, but still keeps the stream up for some time 1-2 seconds (more or less, depends), time which does not let the user to add a new message.
@westey-m do you have some insights about current subject? sample of code would help