fix: preserve streaming usage after AbortSignal cancellation

.changeset/abort-stream-usage.md

@@ -0,0 +1,6 @@
+---
+'@openai/agents-core': patch
+'@openai/agents-openai': patch
+---
+
+fix: preserve streaming usage snapshots across aborts and response_done replacement (#995)

packages/agents-core/src/run.ts

@@ -59,6 +59,7 @@ import { processModelResponseAsync } from './runner/modelOutputs';
 import {
   addStepToRunResult,
   streamStepItemsToRunResult,
+  getUsageSnapshotFromStreamEvent,
   isAbortError,
 } from './runner/streaming';
 import {
@@ -1149,6 +1150,7 @@ export class Runner extends RunHooks<any, AgentOutputType<unknown>> {
           guardrailTracker.throwIfError();
 
           let finalResponse: ModelResponse | undefined = undefined;
+          let latestStreamingUsageSnapshot: Usage | undefined = undefined;
           let inputMarked = false;
           const markInputOnce = () => {
             if (inputMarked || !serverConversationTracker) {
@@ -1168,6 +1170,13 @@ export class Runner extends RunHooks<any, AgentOutputType<unknown>> {
             );
             inputMarked = true;
           };
+          const applyStreamingUsageSnapshot = (usageSnapshot: Usage) => {
+            result.state._context.usage.replaceCurrentRequestSnapshot(
+              usageSnapshot,
+              latestStreamingUsageSnapshot,
+            );
+            latestStreamingUsageSnapshot = usageSnapshot;
+          };
 
           sentInputToModel = true;
           if (!delayStreamInputPersistence) {
@@ -1203,6 +1212,10 @@ export class Runner extends RunHooks<any, AgentOutputType<unknown>> {
             )) {
               guardrailTracker.throwIfError();
               markInputOnce();
+              const usageSnapshot = getUsageSnapshotFromStreamEvent(event);
+              if (usageSnapshot) {
+                applyStreamingUsageSnapshot(usageSnapshot);
+              }
               if (event.type === 'response_done') {
                 const parsed = StreamEventResponseCompleted.parse(event);
                 finalResponse = {
@@ -1211,7 +1224,6 @@ export class Runner extends RunHooks<any, AgentOutputType<unknown>> {
                   responseId: parsed.response.id,
                   requestId: parsed.response.requestId,
                 };
-                result.state._context.usage.add(finalResponse.usage);
               }
               if (result.cancelled) {
                 // When the user's code exits a loop to consume the stream, we need to break

packages/agents-core/src/runner/modelRetry.ts

@@ -89,6 +89,75 @@ function addFailedRetryAttemptsToUsage(
   });
 }
 
+function toUsageSnapshotData(
+  usage: Usage,
+): ConstructorParameters<typeof Usage>[0] {
+  return {
+    requests: usage.requests,
+    inputTokens: usage.inputTokens,
+    outputTokens: usage.outputTokens,
+    totalTokens: usage.totalTokens,
+    inputTokensDetails: usage.inputTokensDetails,
+    outputTokensDetails: usage.outputTokensDetails,
+    ...(usage.requestUsageEntries && usage.requestUsageEntries.length > 0
+      ? {
+          requestUsageEntries: usage.requestUsageEntries.map((entry) => ({
+            inputTokens: entry.inputTokens,
+            outputTokens: entry.outputTokens,
+            totalTokens: entry.totalTokens,
+            inputTokensDetails: entry.inputTokensDetails,
+            outputTokensDetails: entry.outputTokensDetails,
+            ...(entry.endpoint ? { endpoint: entry.endpoint } : {}),
+          })),
+        }
+      : {}),
+  };
+}
+
+function addFailedRetryAttemptsToStreamEvent(
+  event: StreamEvent,
+  failedRetryAttempts: number,
+): StreamEvent {
+  if (failedRetryAttempts <= 0) {
+    return event;
+  }
+
+  if (event.type === 'response_done') {
+    return {
+      ...event,
+      response: {
+        ...event.response,
+        usage: addFailedRetryAttemptsToUsage(
+          new Usage(event.response.usage),
+          failedRetryAttempts,
+        ),
+      },
+    };
+  }
+
+  if (event.type !== 'model') {
+    return event;
+  }
+
+  const usageSnapshot = event.providerData?.usageSnapshot;
+  if (!isRecord(usageSnapshot)) {
+    return event;
+  }
+
+  return {
+    ...event,
+    providerData: {
+      ...(event.providerData ?? {}),
+      usageSnapshot: toUsageSnapshotData(
+        addFailedRetryAttemptsToUsage(
+          new Usage(usageSnapshot as ConstructorParameters<typeof Usage>[0]),
+          failedRetryAttempts,
+        ),
+      ),
+    },
+  };
+}
+
 function withRunnerManagedRetry(request: ModelRequest): ModelRequest {
   return Object.assign({}, request, {
     _internal: {
@@ -736,20 +805,7 @@ export async function* getStreamedResponseWithRetry(
           emittedRawModelEvent = true;
         }
         emittedVisibleEvent = true;
-        if (event.type === 'response_done' && attempt > 1) {
-          yield {
-            ...event,
-            response: {
-              ...event.response,
-              usage: addFailedRetryAttemptsToUsage(
-                new Usage(event.response.usage),
-                attempt - 1,
-              ),
-            },
-          };
-          continue;
-        }
-        yield event;
+        yield addFailedRetryAttemptsToStreamEvent(event, attempt - 1);
       }
       return;
     } catch (error) {

packages/agents-core/src/runner/streaming.ts

@@ -1,4 +1,5 @@
 import logger from '../logger';
+import { Usage } from '../usage';
 import { RunItemStreamEvent, RunItemStreamEventName } from '../events';
 import {
   RunHandoffCallItem,
@@ -12,6 +13,8 @@ import {
   RunToolSearchCallItem,
   RunToolSearchOutputItem,
 } from '../items';
+import type { StreamEvent } from '../types/protocol';
+import { StreamEventResponseCompleted } from '../types/protocol';
 import { StreamedRunResult } from '../result';
 
 export const isAbortError = (error: unknown): boolean => {
@@ -33,6 +36,26 @@ export const isAbortError = (error: unknown): boolean => {
   return false;
 };
 
+export function getUsageSnapshotFromStreamEvent(
+  event: StreamEvent,
+): Usage | undefined {
+  if (event.type === 'response_done') {
+    return new Usage(StreamEventResponseCompleted.parse(event).response.usage);
+  }
+
+  if (
+    event.type !== 'model' ||
+    !event.providerData?.usageSnapshot ||
+    typeof event.providerData.usageSnapshot !== 'object'
+  ) {
+    return undefined;
+  }
+
+  return new Usage(
+    event.providerData.usageSnapshot as ConstructorParameters<typeof Usage>[0],
+  );
+}
+
 function getRunItemStreamEventName(
   item: RunItem,
 ): RunItemStreamEventName | undefined {

packages/agents-core/src/usage.ts

@@ -216,6 +216,92 @@ export class Usage {
       );
     }
   }
+
+  /**
+   * Replaces the latest in-flight request usage snapshot with a newer snapshot.
+   *
+   * This is used for streaming providers that surface provisional usage before
+   * emitting a terminal response event.
+   */
+  replaceCurrentRequestSnapshot(nextUsage: Usage, previousUsage?: Usage) {
+    if (!previousUsage) {
+      this.add(nextUsage);
+      return;
+    }
+
+    this.requests += nextUsage.requests - previousUsage.requests;
+    this.inputTokens += nextUsage.inputTokens - previousUsage.inputTokens;
+    this.outputTokens += nextUsage.outputTokens - previousUsage.outputTokens;
+    this.totalTokens += nextUsage.totalTokens - previousUsage.totalTokens;
+
+    this.#replaceLatestDetails(
+      this.inputTokensDetails,
+      previousUsage.inputTokensDetails,
+      nextUsage.inputTokensDetails,
+    );
+    this.#replaceLatestDetails(
+      this.outputTokensDetails,
+      previousUsage.outputTokensDetails,
+      nextUsage.outputTokensDetails,
+    );
+
+    this.#replaceLatestRequestUsageEntries(
+      Usage.#getTrackedRequestUsageEntries(previousUsage),
+      Usage.#getTrackedRequestUsageEntries(nextUsage),
+    );
+  }
+
+  static #getTrackedRequestUsageEntries(usage: Usage): RequestUsage[] {
+    if (usage.requestUsageEntries?.length) {
+      return usage.requestUsageEntries.map((entry) => new RequestUsage(entry));
+    }
+
+    if (usage.requests === 1 && usage.totalTokens > 0) {
+      return [
+        new RequestUsage({
+          inputTokens: usage.inputTokens,
+          outputTokens: usage.outputTokens,
+          totalTokens: usage.totalTokens,
+          inputTokensDetails: usage.inputTokensDetails[0],
+          outputTokensDetails: usage.outputTokensDetails[0],
+        }),
+      ];
+    }
+
+    return [];
+  }
+
+  #replaceLatestDetails(
+    details: Array<Record<string, number>>,
+    previous: Array<Record<string, number>>,
+    next: Array<Record<string, number>>,
+  ) {
+    details.splice(
+      Math.max(details.length - previous.length, 0),
+      previous.length,
+      ...next,
+    );
+  }
+
+  #replaceLatestRequestUsageEntries(
+    previous: RequestUsage[],
+    next: RequestUsage[],
+  ) {
+    if (previous.length > 0) {
+      const retainedEntries =
+        this.requestUsageEntries?.slice(
+          0,
+          Math.max(this.requestUsageEntries.length - previous.length, 0),
+        ) ?? [];
+      this.requestUsageEntries =
+        retainedEntries.length > 0 ? retainedEntries : undefined;
+    }
+
+    if (next.length > 0) {
+      this.requestUsageEntries ??= [];
+      this.requestUsageEntries.push(...next);
+    }
+  }
 }
 
 export { RequestUsageData, UsageData };