fix(ai): Call onFinish when message ui stream is cancelled in toUIMessageStream - #7902

Anything blocking here? Also hitting this issue and would love if this was merged.

I am not sure this PR solves the problem and it's a bit difficult to test.

The problem here is that ai sdk currently uses a type: "abort" chunk to detect aborts but it looks like this chunk cannot be sent in case of a client abort because the stream can no longer be written to after the abort happens.

Fixing this may require a big change, updating how aborts are detected in a lot of places

I found a little documented field on TransformStream that should fix the problem @lgrammel

The tests I added fail without the cancel fix

I want to understand why this is a problem to begin with. I changed the abort behavior in streamText to not abort a stream any more; instead abort gets turned into regular stream behavior with an abort chunk to prevent special cases like this one.

Can you explain or even better share reproduction code for how such a situation can even occur?

The problem was not strictly related to the abort controller but interrupting the async generator.

While writing this comment I also found another issue which was related to consuming the stream with await for and doing an early break.

I updated and added a new test file with the reproduction which do not pass without the fix and use the full streamText function to make sure they reproduce the real problem.

Basically the smallest repro is this code:

const result = streamText({
  model,
  prompt: 'Generate text',
});

const uiStream = result.toUIMessageStream({
  onFinish: onFinishCallback,
  generateMessageId: () => 'msg-async-iter',
});

// Consume only first few chunks using async iteration, then break
let chunkCount = 0;

for await (const chunk of uiStream) {
  chunkCount++;

  if (chunkCount >= 5) {
    break; // This was not triggering the stream flush(). Instead it just stops consuming the stream and onfinish is never called
  }
}
// onFinish is never called

This code will cause the generator to stop and the stream to stop being consumed, the problem was in createAsyncIterableStream which does not pass a return method in the async iterator. return is necessary to cleanup resources when a await for loop breaks

I updated the code to cancel the stream when the async generator breaks. With the added cancel method in the TransformStream onFinish will be called.

Adding cancel in TransofrmStream now also handles the case when the user cancels the UIMessageStreamResponse, which is another problem that some users could face, some frameworks probably cancel responses when the request abort controller aborts to prevent memory leaks. Smallest repro is this code:

const res = result.toUIMessageStreamResponse()
res.body?.cancel()

I fixed the conflicts with upstream, ready to review

streamText has an onFinish and an onAbort callback. The onFinish callback receives usage information that is not available / incorrect in the case of abort. Therefore only onAbort should be called.

Is there a bug where onAbort is not called? onFinish is not called by design atm.

(this is slightly different for the ui message stream onFinish so there may be more here, just reviewed the first test part)

Are you referring to the test I added in the last commit right?

I added this test to make sure I was understanding when StreamText onFinish runs and this was the wanted behaviour. Notice this test is not about aborts but errors. I tried running this same test in main and it still passes so it is unrelated to my changes.

Do you want me to remove that test?

@remorses maybe i got confused by that test. if it is valuable, would it be possible to move it to a separate pr and integrate it into stream-text.test.ts?

Sure. I can confirm onFinish is still not called during aborts, I can add test for it

@remorses ideally move as much as possible into separate prs. overload is real so i appreciate minimal prs focussed on 1 problem.

@remorses could you look into the merge conflict? Sorry for that, we can look into it ourselves but it might take a bit longer

Fixed the merge conflict