Cost tracking

[ngoiyaeric]

https://docs.litellm.ai/docs/completion/token_usage

Implement cost tracking for tool calling and all peripherals.

[ngoiyaeric]

@coderabbitai plan use litellm documentation to make this extensible with tools

[coderabbitai]

Coding Plan

Summary

• Treat the LiteLLM link as conceptual guidance and implement cost tracking natively on the existing Vercel AI SDK, capturing token usage via onFinish/returned usage across all six agents.
• Introduce a central cost configuration (per-model token rates + per-tool flat costs) and a cost-calculation utility, with getModel() extended to expose the active model identifier for pricing lookups.
• Persist every LLM and peripheral tool call as a durable usage_events row via a single best-effort recordUsageEvent action, with an aggregation read action for the UI.
• Replace the existing hardcoded usage/credits UI with real per-user cost and token data; deliberately exclude Stripe billing and credit enforcement as out of scope.

Design Choices

Design Choice 1: Should this ticket implement full credit enforcement/billing, or only cost tracking?

Options Considered:

1. Implement tracking + persistence + display only (record and show costs).
2. Implement tracking plus credit deduction/enforcement and Stripe integration.

Chosen Option: 1

Rationale: The ticket scope is "cost tracking"; building enforcement/billing would be scope creep. We capture, persist, and display cost data, leaving the existing Stripe link untouched.

Design Choice 2: How should token usage be captured from streaming agents?

Options Considered:

1. Use the onFinish callback (provides usage, finishReason) on streaming calls and read returned usage on non-streaming calls.
2. Manually await result.usage after streaming completes.

Chosen Option: 1

Rationale: onFinish is the SDK-native hook that fires exactly once with final usage for both streaming and tool-augmented runs, and avoids restructuring the existing stream-consumption loops.

Design Choice 3: How is "tool calling and peripherals" cost determined when external APIs (Tavily, Jina, Composio/Mapbox, Serper) return no token usage?

Options Considered:

1. Assign each tool a configurable per-invocation cost (flat rate) recorded as a usage event.
2. Attempt to parse provider-specific billing responses per tool.

Chosen Option: 1

Rationale: A central per-tool cost table is simple, consistent with the LLM pricing config, and adjustable without code changes; per-provider billing parsing is unreliable and out of scope.

💡 User Tips

Regenerate the plan with different choices with @coderabbitai <feedback>.

Implementation Steps

Phase 1: Cost & Usage Data Foundation

Establish the pricing knowledge, persistence schema, and shared types that all later phases depend on. This phase introduces the concept of a "usage event" (a single billable unit — either an LLM call or a peripheral tool call) and the per-user aggregation needed by the UI.

Task 1: Add pricing configuration and cost-calculation utility

Create a single source of truth mapping models and tools to costs, plus a helper that converts token counts into dollar cost.

• Add a new config module under lib/costs/ defining per-model input/output token rates (keyed by the model identifier returned by the AI SDK model object's modelId, covering the GPT-4o, Gemini, Bedrock Claude, and Grok models used in getModel()).
• In the same module define a per-tool flat cost table for peripherals (search/Tavily, retrieve/Jina, geospatialQueryTool/Composio-Mapbox, videoSearch/Serper).
• Export a calculateLlmCost({ modelId, promptTokens, completionTokens }) function and a getToolCost(toolName) function. Fall back to a configurable default rate when a model id is unmatched.
• Update getModel() in lib/utils/index.ts to additionally expose the selected model's identifier (or return it alongside the model) so callers can look up pricing without duplicating the selection logic.

Task 2: Add usage/cost database tables and migration

Introduce durable storage for individual usage events and a per-user rolling aggregate.

• Extend lib/db/schema.ts with a usage_events table: id, user_id (FK to users), chat_id (nullable FK to chats), kind ('llm' | 'tool'), source (model id or tool name), prompt_tokens, completion_tokens, total_tokens (nullable for tool events), cost (numeric/decimal USD), created_at.
• Optionally add a user_usage_summary table (or compute aggregates on read) holding user_id, total_cost, total_tokens, updated_at for fast UI reads; prefer on-read aggregation first to minimize new write paths.
• Generate the migration via drizzle-kit generate (producing the next 000N_<slug>.sql plus updated _journal.json and snapshot), following the existing migration convention.

Task 3: Add TypeScript types for usage and cost

Provide shared types consumed by agents, actions, and UI.

• Extend lib/types/index.ts with types such as UsageEvent, LlmUsage (promptTokens/completionTokens/totalTokens), and UsageSummary (aggregate cost/tokens) used across server actions and the UI.

🤖 Prompt for AI agents

Implement the cost & usage data foundation — the pricing config, database
schema, and shared TypeScript types that all subsequent phases will depend on.

**1. Create `lib/costs/` module:**
- Create a new file (e.g., `lib/costs/index.ts`) that defines:
- A per-model pricing map keyed by `modelId` (as returned by the AI SDK),
covering GPT-4o, Gemini, Bedrock Claude, and Grok models used in `getModel()`.
Each entry should include input and output token rates (USD per token).
- A per-tool flat cost table for: `search` (Tavily), `retrieve` (Jina),
`geospatialQueryTool` (Composio-Mapbox), and `videoSearch` (Serper).
- Export a `calculateLlmCost({ modelId, promptTokens, completionTokens })`
function that returns a dollar cost, with a fallback default rate for unmatched
model IDs.
- Export a `getToolCost(toolName)` function that returns the flat per-invocation
cost for a given tool name.

**2. Update `lib/utils/index.ts`:**
- Modify `getModel()` to also expose the selected model's `modelId` identifier
(either by returning it alongside the model object, or by another accessor
pattern), so callers can look up pricing without duplicating model selection
logic.

**3. Extend `lib/db/schema.ts`:**
- Add a `usage_events` table with columns: `id`, `user_id` (FK to `users`),
`chat_id` (nullable FK to `chats`), `kind` (enum: `'llm'` | `'tool'`), `source`
(model id or tool name string), `prompt_tokens`, `completion_tokens`,
`total_tokens` (nullable, for tool events), `cost` (numeric/decimal USD),
`created_at`.
- Optionally add a `user_usage_summary` table with `user_id`, `total_cost`,
`total_tokens`, `updated_at`; however, prefer on-read aggregation first to
minimize new write paths.

**4. Generate the database migration:**
- Run `drizzle-kit generate` to produce the next migration file
(`000N_<slug>.sql`), updated `_journal.json`, and snapshot, following the
existing convention in `drizzle/migrations`.

**5. Extend `lib/types/index.ts`:**
- Add `UsageEvent` type matching the `usage_events` schema.
- Add `LlmUsage` type with fields: `promptTokens`, `completionTokens`,
`totalTokens`.
- Add `UsageSummary` type for aggregate cost/token totals used by the UI.

Phase 2: Instrument AI Calls and Tools

Capture real consumption data at the point each cost is incurred. LLM token usage is read from AI SDK callbacks; peripheral tool costs are recorded inside each tool's execute.

Task 1: Capture token usage from all agents

Wire usage capture into every AI SDK invocation and translate it into a recorded cost event.

• Add onFinish callbacks to the streaming agents (lib/agents/researcher.tsx, lib/agents/writer.tsx, lib/agents/inquire.tsx) that receive usage and forward it (with the active model id, current user, and chat id) to the persistence action from Phase 3.
• For non-streaming agents (lib/agents/task-manager.tsx, lib/agents/report/executive-summary.ts, lib/agents/report/strategic-synthesis.ts), destructure usage from the generateText/generateObject result and forward it the same way.
• Ensure user/chat context (already available in app/actions.tsx orchestration) is threaded into agents so events can be attributed correctly.

Task 2: Record peripheral/tool costs

Track each external tool invocation as a usage event.

• In each tool's execute (lib/agents/tools/search.tsx, retrieve.tsx, geospatial.tsx, video-search.tsx), after a successful external call, record a 'tool'-kind usage event using getToolCost(toolName).
• Prefer a thin shared wrapper or helper (e.g., in lib/costs/) invoked from each tool to avoid duplicating recording logic; pass through user/chat context provided via getTools() options in lib/agents/tools/index.tsx.

🤖 Prompt for AI agents

Instrument all AI agent calls and peripheral tool executions to capture real
usage data at the point each cost is incurred.

**1. Instrument streaming agents with `onFinish` callbacks:**
- In `lib/agents/researcher.tsx`, `lib/agents/writer.tsx`, and
`lib/agents/inquire.tsx`, add an `onFinish` callback to each
`streamText`/`streamObject` invocation.
- The `onFinish` handler receives `usage` (token counts) and should forward it,
along with the active `modelId`, current `userId`, and `chatId`, to the
`recordUsageEvent` persistence action (from Phase 3).

**2. Instrument non-streaming agents:**
- In `lib/agents/task-manager.tsx`, `lib/agents/report/executive-summary.ts`,
and `lib/agents/report/strategic-synthesis.ts`, destructure `usage` from the
returned result of `generateText`/`generateObject`.
- Forward `usage`, `modelId`, `userId`, and `chatId` to `recordUsageEvent` the
same way as streaming agents.

**3. Thread user/chat context into agents:**
- Ensure `userId` and `chatId` (already available in `app/actions.tsx`
orchestration) are passed down into each agent so usage events can be correctly
attributed.

**4. Record peripheral/tool costs:**
- In each tool's `execute` function (`lib/agents/tools/search.tsx`,
`retrieve.tsx`, `geospatial.tsx`, `video-search.tsx`), after a successful
external API call, invoke `getToolCost(toolName)` and record a `'tool'`-kind
usage event.
- Create a thin shared helper (e.g., in `lib/costs/`) to encapsulate the
recording logic, avoiding duplication across tools.
- Ensure `userId` and `chatId` context is available to tools by threading it
through `getTools()` options in `lib/agents/tools/index.tsx`.

Phase 3: Persistence and Aggregation

Provide the server-side write and read paths that connect instrumentation to the database and the UI.

Task 1: Add usage-recording server action

Centralize how usage events are written so agents and tools share one path.

• Create lib/actions/usage.ts exporting recordUsageEvent(event) that inserts into usage_events using the shared db client; for LLM events it computes cost via calculateLlmCost, for tool events via getToolCost.
• Make the write resilient (best-effort, non-blocking) so a tracking failure never breaks the chat/AI flow.

Task 2: Add usage-aggregation read action

Expose per-user totals for the UI.

• In lib/actions/usage.ts, add getUserUsageSummary(userId) returning total cost, total tokens, and a list of recent events (for the transactions table), via Drizzle aggregate queries over usage_events.

🤖 Prompt for AI agents

Create the server-side persistence and aggregation layer in
`lib/actions/usage.ts` that connects instrumentation to the database and the UI.

**1. Create `lib/actions/usage.ts`:**
- Export `recordUsageEvent(event)` that:
  - Accepts a usage event payload (typed using `UsageEvent` from Phase 1).
- For `'llm'`-kind events, computes the dollar cost via `calculateLlmCost({
modelId, promptTokens, completionTokens })`.
  - For `'tool'`-kind events, computes cost via `getToolCost(toolName)`.
- Inserts the complete event row into `usage_events` using the shared `db`
client (from `lib/db/index.ts`).
- Is **best-effort and non-blocking**: wrap the insert in a try/catch (or
`.catch()`) so any failure is silently logged and never propagates to break the
chat/AI flow.

**2. Add `getUserUsageSummary(userId)` in the same file:**
- Accepts a `userId` and queries the `usage_events` table via Drizzle aggregate
functions.
- Returns: total cost (sum of `cost`), total tokens (sum of `total_tokens`), and
a list of recent usage events for display in the transactions table.
- Return shape should conform to the `UsageSummary` type defined in Phase 1.

Phase 4: Surface Real Usage in the UI

Replace the hardcoded placeholders with live data so cost tracking is visible to users.

Task 1: Wire the usage dashboard to real data

Render actual usage and cost in the main usage panel.

• Update components/usage-view.tsx to fetch from getUserUsageSummary (via a server action call or passed-in props) instead of useState(0) and the three hardcoded transactions, displaying real total cost/tokens and recent usage events.

Task 2: Wire the sidebar credits preview

Reflect real usage in the collapsible sidebar widget.

• Update components/sidebar/chat-history-client.tsx "Credits Preview" to consume the same usage summary, replacing the hardcoded 0 value and 0% progress bar.

🤖 Prompt for AI agents

Replace the hardcoded usage/credits placeholders in the UI with live data
sourced from the `getUserUsageSummary` server action.

**1. Update `components/usage-view.tsx`:**
- Remove the `useState(0)` placeholder and the three hardcoded transaction
entries.
- Fetch real usage data by calling `getUserUsageSummary(userId)` (either as a
server component, via a server action call, or via passed-in props).
- Display real values for: total cost, total tokens, and a list of recent usage
events in the transactions table.

**2. Update `components/sidebar/chat-history-client.tsx`:**
- In the "Credits Preview" widget, replace the hardcoded `0` value and `0%`
progress bar with values sourced from the same `getUserUsageSummary` usage
summary.
- Ensure the same data fetch/prop pattern is used consistently with the usage
dashboard update above.

Research

This is a Next.js app using the Vercel AI SDK (not LiteLLM, despite the ticket link) with multiple provider SDKs (@ai-sdk/openai, @ai-sdk/google, @ai-sdk/amazon-bedrock, @ai-sdk/xai). Model selection is centralized in getModel() (lib/utils/index.ts). AI work is split across six agents (researcher, writer, task-manager, inquire, executive-summary, strategic-synthesis), none of which currently read token usage. Tools follow a uniform shape and are assembled in lib/agents/tools/index.tsx. Persistence uses Drizzle ORM over PostgreSQL (schema in lib/db/schema.ts, shared client in lib/db/index.ts, migrations in drizzle/migrations). There is no existing usage/cost/billing data model — the usage UI (components/usage-view.tsx) and sidebar "Credits Preview" are hardcoded placeholders.

---

🚀 Next Steps

🤖 All AI agent prompts combined

Task: 1

Implement the cost & usage data foundation — the pricing config, database
schema, and shared TypeScript types that all subsequent phases will depend on.

**1. Create `lib/costs/` module:**
- Create a new file (e.g., `lib/costs/index.ts`) that defines:
- A per-model pricing map keyed by `modelId` (as returned by the AI SDK),
covering GPT-4o, Gemini, Bedrock Claude, and Grok models used in `getModel()`.
Each entry should include input and output token rates (USD per token).
- A per-tool flat cost table for: `search` (Tavily), `retrieve` (Jina),
`geospatialQueryTool` (Composio-Mapbox), and `videoSearch` (Serper).
- Export a `calculateLlmCost({ modelId, promptTokens, completionTokens })`
function that returns a dollar cost, with a fallback default rate for unmatched
model IDs.
- Export a `getToolCost(toolName)` function that returns the flat per-invocation
cost for a given tool name.

**2. Update `lib/utils/index.ts`:**
- Modify `getModel()` to also expose the selected model's `modelId` identifier
(either by returning it alongside the model object, or by another accessor
pattern), so callers can look up pricing without duplicating model selection
logic.

**3. Extend `lib/db/schema.ts`:**
- Add a `usage_events` table with columns: `id`, `user_id` (FK to `users`),
`chat_id` (nullable FK to `chats`), `kind` (enum: `'llm'` | `'tool'`), `source`
(model id or tool name string), `prompt_tokens`, `completion_tokens`,
`total_tokens` (nullable, for tool events), `cost` (numeric/decimal USD),
`created_at`.
- Optionally add a `user_usage_summary` table with `user_id`, `total_cost`,
`total_tokens`, `updated_at`; however, prefer on-read aggregation first to
minimize new write paths.

**4. Generate the database migration:**
- Run `drizzle-kit generate` to produce the next migration file
(`000N_<slug>.sql`), updated `_journal.json`, and snapshot, following the
existing convention in `drizzle/migrations`.

**5. Extend `lib/types/index.ts`:**
- Add `UsageEvent` type matching the `usage_events` schema.
- Add `LlmUsage` type with fields: `promptTokens`, `completionTokens`,
`totalTokens`.
- Add `UsageSummary` type for aggregate cost/token totals used by the UI.
===============================================================================

Task: 2

Instrument all AI agent calls and peripheral tool executions to capture real
usage data at the point each cost is incurred.

**1. Instrument streaming agents with `onFinish` callbacks:**
- In `lib/agents/researcher.tsx`, `lib/agents/writer.tsx`, and
`lib/agents/inquire.tsx`, add an `onFinish` callback to each
`streamText`/`streamObject` invocation.
- The `onFinish` handler receives `usage` (token counts) and should forward it,
along with the active `modelId`, current `userId`, and `chatId`, to the
`recordUsageEvent` persistence action (from Phase 3).

**2. Instrument non-streaming agents:**
- In `lib/agents/task-manager.tsx`, `lib/agents/report/executive-summary.ts`,
and `lib/agents/report/strategic-synthesis.ts`, destructure `usage` from the
returned result of `generateText`/`generateObject`.
- Forward `usage`, `modelId`, `userId`, and `chatId` to `recordUsageEvent` the
same way as streaming agents.

**3. Thread user/chat context into agents:**
- Ensure `userId` and `chatId` (already available in `app/actions.tsx`
orchestration) are passed down into each agent so usage events can be correctly
attributed.

**4. Record peripheral/tool costs:**
- In each tool's `execute` function (`lib/agents/tools/search.tsx`,
`retrieve.tsx`, `geospatial.tsx`, `video-search.tsx`), after a successful
external API call, invoke `getToolCost(toolName)` and record a `'tool'`-kind
usage event.
- Create a thin shared helper (e.g., in `lib/costs/`) to encapsulate the
recording logic, avoiding duplication across tools.
- Ensure `userId` and `chatId` context is available to tools by threading it
through `getTools()` options in `lib/agents/tools/index.tsx`.
===============================================================================

Task: 3

Create the server-side persistence and aggregation layer in
`lib/actions/usage.ts` that connects instrumentation to the database and the UI.

**1. Create `lib/actions/usage.ts`:**
- Export `recordUsageEvent(event)` that:
  - Accepts a usage event payload (typed using `UsageEvent` from Phase 1).
- For `'llm'`-kind events, computes the dollar cost via `calculateLlmCost({
modelId, promptTokens, completionTokens })`.
  - For `'tool'`-kind events, computes cost via `getToolCost(toolName)`.
- Inserts the complete event row into `usage_events` using the shared `db`
client (from `lib/db/index.ts`).
- Is **best-effort and non-blocking**: wrap the insert in a try/catch (or
`.catch()`) so any failure is silently logged and never propagates to break the
chat/AI flow.

**2. Add `getUserUsageSummary(userId)` in the same file:**
- Accepts a `userId` and queries the `usage_events` table via Drizzle aggregate
functions.
- Returns: total cost (sum of `cost`), total tokens (sum of `total_tokens`), and
a list of recent usage events for display in the transactions table.
- Return shape should conform to the `UsageSummary` type defined in Phase 1.
===============================================================================

Task: 4

Replace the hardcoded usage/credits placeholders in the UI with live data
sourced from the `getUserUsageSummary` server action.

**1. Update `components/usage-view.tsx`:**
- Remove the `useState(0)` placeholder and the three hardcoded transaction
entries.
- Fetch real usage data by calling `getUserUsageSummary(userId)` (either as a
server component, via a server action call, or via passed-in props).
- Display real values for: total cost, total tokens, and a list of recent usage
events in the transactions table.

**2. Update `components/sidebar/chat-history-client.tsx`:**
- In the "Credits Preview" widget, replace the hardcoded `0` value and `0%`
progress bar with values sourced from the same `getUserUsageSummary` usage
summary.
- Ensure the same data fetch/prop pattern is used consistently with the usage
dashboard update above.

💡 Iterate on the plan with: `@coderabbitai` <feedback>

Example Feedback
- `@coderabbitai` You can skip phase 3. Add a simple unit test case for phase 2.
- `@coderabbitai` For design choice 1 go ahead with option 3 and replan.

---

💬 Have feedback or questions? Drop into our discord!