Cost tracking #698

Activity

1. ngoiyaeric commented on Jun 27, 2026

   

   ngoiyaeric

   on Jun 27, 2026

   CollaboratorAuthor

   More actions

   @coderabbitai plan use litellm documentation to make this extensible with tools
2. coderabbitai commented on Jun 27, 2026

   

   coderabbitaibot

   on Jun 27, 2026 – with coderabbitai

   Contributor

   More actions

   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!