Skip to content

Conversation

@feelfetch
Copy link

Summary

This PR adds token usage tracking to the @ai-sdk-tools/agents package, addressing #57. Previously, toUIMessageStream() did not expose token usage data, making it impossible to track consumption, calculate costs, or display usage metrics.

Features

1. Usage in onFinish Callback

The onFinish callback now includes accumulated token usage across all agent rounds:

const response = await agent.toUIMessageStream({
  message,
  onFinish: async (event) => {
    console.log(event.usage);
    // { inputTokens: 150, outputTokens: 200, totalTokens: 350 }
    
    // Calculate costs, save to database, enforce limits, etc.
  },
});

2. Streamed Usage Data Part

Usage is streamed as a transient data-usage data part, available on the client via useDataPart:

import { useDataPart } from "@ai-sdk-tools/store";

const [usageData] = useDataPart<{ usage: LanguageModelUsage }>("usage");
// Access usageData.usage.inputTokens, outputTokens, totalTokens

3. OpenTelemetry Support

Added experimental_telemetry option for observability platform integration:

const response = await agent.toUIMessageStream({
  message,
  experimental_telemetry: {
    isEnabled: true,
    functionId: "my-agent",
  },
});

Use Cases Enabled

  • ✅ Display token counts to users
  • ✅ Calculate costs with tools like TokenLens
  • ✅ Track context window usage
  • ✅ Monitor usage patterns for optimization
  • ✅ Implement per-user/organization token quotas
  • ✅ Persist usage history to database
  • ✅ OpenTelemetry tracing for observability

Implementation Details

  • Usage is accumulated across all agent rounds (orchestrator + specialists)
  • Uses AI SDK v5's result.totalUsage for multi-step generations
  • Gracefully handles cases where context is undefined
  • No breaking changes - fully backward compatible

New Exports

// Types
export type { AgentStreamOnFinishCallback, AgentStreamOnFinishEvent } from '@ai-sdk-tools/agents';

// Helper
export { writeUsage } from '@ai-sdk-tools/agents';

Example Integration

The example app now includes a TokenUsage component in the chat header that displays:

  • Context window usage percentage
  • Input/output token breakdown
  • Estimated costs via TokenLens
CleanShot 2025-11-29 at 16 24 45

@vercel
Copy link

vercel bot commented Nov 29, 2025

@feelfetch is attempting to deploy a commit to the Pontus Abrahamsson's projects Team on Vercel.

A member of the Team first needs to authorize it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant