fix(extensions): handle NaN usage tokens

Author: dkundel-openai

packages/agents-extensions/src/aiSdk.ts

@@ -443,10 +443,19 @@ export class AiSdkModel implements Model {
         return {
           responseId: result.response?.id ?? 'FAKE_ID',
           usage: new Usage({
-            inputTokens: result.usage.promptTokens,
-            outputTokens: result.usage.completionTokens,
+            inputTokens: Number.isNaN(result.usage?.promptTokens)
+              ? 0
+              : (result.usage?.promptTokens ?? 0),
+            outputTokens: Number.isNaN(result.usage?.completionTokens)
+              ? 0
+              : (result.usage?.completionTokens ?? 0),
             totalTokens:
-              result.usage.promptTokens + result.usage.completionTokens,
+              (Number.isNaN(result.usage?.promptTokens)
+                ? 0
+                : (result.usage?.promptTokens ?? 0)) +
+              (Number.isNaN(result.usage?.completionTokens)
+                ? 0
+                : (result.usage?.completionTokens ?? 0)),
           }),
           output,
         };
@@ -608,8 +617,12 @@ export class AiSdkModel implements Model {
             break;
           }
           case 'finish': {
-            usagePromptTokens = part.usage.promptTokens;
-            usageCompletionTokens = part.usage.completionTokens;
+            usagePromptTokens = Number.isNaN(part.usage?.promptTokens)
+              ? 0
+              : (part.usage?.promptTokens ?? 0);
+            usageCompletionTokens = Number.isNaN(part.usage?.completionTokens)
+              ? 0
+              : (part.usage?.completionTokens ?? 0);
             break;
           }
           case 'error': {

packages/agents-extensions/test/aiSdk.test.ts

@@ -468,6 +468,42 @@ describe('AiSdkModel.getResponse', () => {
       content: 'inst',
     });
   });
+
+  test('handles NaN usage in doGenerate', async () => {
+    const model = new AiSdkModel(
+      stubModel({
+        async doGenerate() {
+          return {
+            text: '',
+            finishReason: 'stop',
+            usage: { promptTokens: Number.NaN, completionTokens: Number.NaN },
+            providerMetadata: {},
+            rawCall: { rawPrompt: '', rawSettings: {} },
+          };
+        },
+      }),
+    );
+
+    const res = await withTrace('t', () =>
+      model.getResponse({
+        input: 'hi',
+        tools: [],
+        handoffs: [],
+        modelSettings: {},
+        outputType: 'text',
+        tracing: false,
+      } as any),
+    );
+
+    expect(res.usage).toEqual({
+      requests: 1,
+      inputTokens: 0,
+      outputTokens: 0,
+      totalTokens: 0,
+      inputTokensDetails: [],
+      outputTokensDetails: [],
+    });
+  });
 });
 
 describe('AiSdkModel.getStreamedResponse', () => {
@@ -639,6 +675,43 @@ describe('AiSdkModel.getStreamedResponse', () => {
       content: 'inst',
     });
   });
+
+  test('handles NaN usage in stream finish event', async () => {
+    const parts = [
+      { type: 'text-delta', textDelta: 'a' },
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        usage: { promptTokens: Number.NaN, completionTokens: Number.NaN },
+      },
+    ];
+    const model = new AiSdkModel(
+      stubModel({
+        async doStream() {
+          return {
+            stream: partsStream(parts),
+            rawCall: { rawPrompt: '', rawSettings: {} },
+          } as any;
+        },
+      }),
+    );
+
+    let final: any;
+    for await (const ev of model.getStreamedResponse({
+      input: 'hi',
+      tools: [],
+      handoffs: [],
+      modelSettings: {},
+      outputType: 'text',
+      tracing: false,
+    } as any)) {
+      if (ev.type === 'response_done') {
+        final = ev.response.usage;
+      }
+    }
+
+    expect(final).toEqual({ inputTokens: 0, outputTokens: 0, totalTokens: 0 });
+  });
 });
 
 describe('AiSdkModel', () => {