Skip to content

feat(ai): expose inputMessages and partialText in onAbort callback#14141

Open
Hovo-Dev wants to merge 5 commits intovercel:mainfrom
Hovo-Dev:fix/stream-text-abort-usage
Open

feat(ai): expose inputMessages and partialText in onAbort callback#14141
Hovo-Dev wants to merge 5 commits intovercel:mainfrom
Hovo-Dev:fix/stream-text-abort-usage

Conversation

@Hovo-Dev
Copy link
Copy Markdown

@Hovo-Dev Hovo-Dev commented Apr 5, 2026

Background

When a streamText call is aborted — via AbortSignal, the user stopping generation, or a dropped connection — the onAbort callback fires but receives no usage data. onFinish never fires on abort at all. This makes it impossible to track token consumption for billing or quota purposes when streams don't complete naturally.

Summary

packages/ai/src/generate-text/stream-text.ts

  • Added usage?: LanguageModelUsage and totalUsage?: LanguageModelUsage to StreamTextOnAbortCallback
  • Added inputMessages: ModelMessage[] — the full message list passed to the model for the aborted step (initial messages + completed prior step tool
    call/result pairs)
  • Added partialText?: string — the text actively being streamed at abort time; undefined if no text was generated; resets at each step boundary
  • Introduced currentStepUsage, completedStepsUsage, currentStepInputMessages, and pullLevelTextContent as outer-scope tracking variables
  • Wired all fields into the existing abort() function — if the provider had not yet sent usage data before the abort, usage is undefined (honest, no
    estimation)

packages/openai/src/chat/openai-chat-language-model.ts

  • OpenAI's finish chunk (which carries real usage) was only emitted in flush(), which never runs on abort
  • Now emits the finish chunk as soon as usage data arrives in the stream transform(), with a finishSent guard to prevent double-emission in flush()

content/docs/07-reference/01-ai-sdk-core/02-stream-text.mdx U

  • Documented usage, totalUsage, inputMessages, and partialText fields in the onAbort callback parameters

content/docs/03-ai-sdk-core/50-error-handling.mdx

  • Updated onAbort section to describe all four new fields

Usage

onAbort({ totalUsage, usage, inputMessages, partialText }) {                                      
  const completedTokens = totalUsage?.totalTokens ?? 0;                                                                                                     

  // Use real usage if provider sent it, otherwise estimate with your own tokenizer:
  const abortedStepTokens = usage?.totalTokens                         
    ?? (countTokens(inputMessages) + countTokens(partialText ?? ''));                                                                                       
                                                                                                                                                            
  trackUsage(completedTokens + abortedStepTokens);                                                                                                          
}                                                                                                                                                           

Manual Verification

Verified via automated tests that exercise the full stream processing pipeline:

  1. Abort fires before any usage chunk → usage and totalUsage are undefined
  2. Abort fires after the provider sent a usage chunk → usage and totalUsage contain the real token counts
  3. Multi-step stream, abort in step 2 → totalUsage correctly sums step 1 and step 2 usage
  4. Abort before first chunk → inputMessages contains the step's input, partialText is undefined
  5. Abort mid-text → partialText contains the streamed text up to abort
  6. Multi-step abort → inputMessages includes tool call/result pairs from prior steps

All 2200+ existing tests continue to pass.

Checklist

  • Tests have been added / updated (for bug fixes / features)
  • Documentation has been added / updated (for bug fixes / features)
  • A patch changeset for relevant packages has been added (for bug fixes / features - run pnpm changeset in the project root)
  • I have reviewed this pull request (self-review)

Related Issues

Fixes #7628
Fixes #7805

When a streamText call is aborted, the onAbort callback now receives
usage (current step) and totalUsage (all steps combined) containing
whatever real token data the provider had reported before the abort.
If the provider had not yet sent a usage chunk, the values are
undefined — no estimation is used.

The OpenAI provider is updated to emit its finish chunk (which carries
real usage) as soon as usage data arrives in the stream, rather than
waiting for the stream to fully close. This ensures usage is available
even when an abort fires before natural completion.

Fixes vercel#7628
Fixes vercel#7805
@tigent tigent bot added ai/core core functions like generateText, streamText, etc. Provider utils, and provider spec. ai/provider related to a provider package. Must be assigned together with at least one `provider/*` label documentation Improvements or additions to documentation feature New feature or request provider/openai Issues related to the @ai-sdk/openai provider labels Apr 5, 2026
Hovo-Dev added 2 commits April 5, 2026 13:56
The finish chunk was unreachable — OpenAI sends usage in a chunk with
empty choices[], so choice?.delta == null triggered an early return
before the emission block. Move the check before the delta guard.
@Hovo-Dev Hovo-Dev force-pushed the fix/stream-text-abort-usage branch from 2dfe4e0 to de7c3a3 Compare April 5, 2026 10:13
@Hovo-Dev
Copy link
Copy Markdown
Author

Hovo-Dev commented Apr 5, 2026

Hi @gr2m,

Took a look at this and managed to get it working cleanly — happy to share the approach.

The SDK was already holding the token data we needed, it just wasn't being handed off when a stream was aborted. We also noticed that usage from OpenAI arrives mid-stream but was only forwarded at natural close, so we made sure it gets captured as soon as it arrives. That way, whether the stream finishes or gets cut short, the data is there.

Looking forward to your feedback!

@Hovo-Dev Hovo-Dev changed the title feat(ai): expose usage metrics in onAbort callback feat(ai): expose inputMessages and partialText in onAbort callback Apr 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ai/core core functions like generateText, streamText, etc. Provider utils, and provider spec. ai/provider related to a provider package. Must be assigned together with at least one `provider/*` label documentation Improvements or additions to documentation feature New feature or request provider/openai Issues related to the @ai-sdk/openai provider

Projects

None yet

Development

Successfully merging this pull request may close these issues.

A way to get Token Usage when a request is Aborted in streamText [Feature Request] Token usage unavailable during streaming abort/interruption

1 participant