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