feat: Per-run cost tracking, phase-level metering, and budget caps

[ryaneggz]

Summary

Add per-run cost and duration tracking to Orchestra. Every LLM invocation will record token counts, cost (USD), duration, and phase metadata. Aggregated metrics are surfaced per thread, with configurable budget caps that integrate with the existing HITL interrupt system. The frontend displays real-time cost badges and a detailed breakdown panel.

Key Deliverables

• Per-turn metrics capture: model, input/output tokens, cost, duration, phase (planning, execution, reflection, tool_call)
• Aggregate metrics: per-thread and per-phase cost summaries
• Provider pricing tables: hardcoded defaults for OpenAI, Anthropic, Google models; admin-configurable overrides via API
• Real-time cost display: running cost badge in chat header, updated via SSE streaming events
• Budget caps: per-thread and per-assistant budget limits with configurable actions (pause via HITL interrupt, warn, stop)
• Historical cost data: per-turn breakdown and cost analytics accessible through API and frontend panel

Implementation Scope

Backend

• New schemas: TurnMetrics, ThreadCostSummary, RunBudget in backend/src/schemas/entities/metrics.py
• Pricing constants in backend/src/constants/pricing.py
• Metrics capture middleware in backend/src/utils/metrics.py
• Budget enforcement service with HITL integration
• REST API endpoints for metrics, pricing CRUD, and budget management
• Alembic migrations for turn_metrics, run_budgets, and pricing_overrides tables

Frontend

• Running cost badge in chat header
• Cost breakdown panel (per-turn table, per-phase summary)
• Budget configuration dialog
• Warning indicators at 75%/100% thresholds

Testing

• Unit tests for pricing calculation, budget thresholds, metrics aggregation
• Integration tests for end-to-end metrics capture and budget enforcement
• Frontend component tests

Spec

Full specification: .claude/specs/cost-tracking.md

Implementation plan: .claude/plans/cost-tracking.md

Milestone

v0.9.0 -- Harness Design

[ryaneggz]

Frontend Implementation Analysis for Cost Tracking

1. Architecture Overview & Key Integration Points

The frontend uses a layered architecture: hooks (useChat, useThread) -> context (ChatContext) -> components (ChatPanel, ChatMessages, ChatNav). All streaming data flows through the useChat hook's handleMessages() function, which dispatches on streamMode (the first element of each SSE payload tuple).

SSE event flow:

• FetchStreamReader / ResponseBodyReader (frontend/src/lib/utils/fetchStreamReader.ts) parses raw SSE lines into typed StreamEvent objects.
• DistributedStreamSource / SyncStreamSource (frontend/src/lib/utils/streamSource.ts) wraps readers behind a unified StreamSource interface.
• useChat.ts lines 260-315 (startManagedStream) converts events via convertEventToLegacy() -> sseHandler() -> handleMessages().
• handleMessages() (line 625) dispatches on streamMode: "metadata", "messages", "values", "error", "aborted", "mcp_sandbox_unreachable".

2. Where the Cost Badge Should Live

Recommended: ChatUtilityRow (primary) + ChatNav header (secondary)

ChatUtilityRow (frontend/src/components/chat/ChatUtilityRow.tsx) is the ideal primary location. This component already renders alongside ChatInput inside ChatComposer and contains the Sandbox selector pill and Tasks counter pill. A cost badge pill would be a natural third element using the identical styling pattern:

h-8 rounded-full border border-border/60 px-3 text-xs text-muted-foreground

This pill-button pattern (icon + label + value) is already used for:

• Sandbox (ThreadSandboxStatus): <SquareTerminal /> Sandbox <value>
• Tasks (TodoList popover): <ListTodo /> Tasks <count>
• Files: <FolderCode /> Files <count>

A CostBadge component would follow the same pattern: <DollarSign /> Cost <$0.12> -- clicking opens the CostBreakdownPanel popover.

ChatNav (frontend/src/components/nav/ChatNav.tsx) is a secondary option for a persistent, always-visible cost indicator (small text or badge next to the Share/NewThread buttons). However, ChatNav is a thin header bar and currently only contains icon buttons (Save, Share, New Thread, Help, ColorMode). Adding a cost display here would break the icon-only pattern. I would recommend keeping the primary interaction in ChatUtilityRow and optionally showing a minimal read-only cost in the loading indicator area.

Per-message cost display: The Message component (frontend/src/components/lists/ChatMessages.tsx, lines 16-248) already displays model and streamingRate (tok/s, token count) on the last AI message (line 222-229). Per-turn cost could be added here inline, after the token count, as a $X.XX suffix on the same muted-foreground text line. This requires each message to carry cost metadata.

3. SSE Integration: How to Add cost_update Events

The spec calls for a cost_update SSE event. Here is the exact integration path:

Step 1: Add cost_update to the stream entity types
In frontend/src/lib/entities/stream.ts, add:

• New CostUpdateEvent type to the SSEEvent union
• New SSEEventType entry "cost_update"

Step 2: Extend the SSE parsers
Both FetchStreamReader.parseEvent() and ResponseBodyReader.parseEvent() in frontend/src/lib/utils/fetchStreamReader.ts have a switch(type) block (lines 159-188, 305-334). Add a case "cost_update" that returns the event.

Step 3: Wire through the stream source
convertEventToLegacy() in useChat.ts (line 204-221) maps StreamEvent to legacy [type, data] tuples. Add a case for "cost_update".

Step 4: Handle in handleMessages()
In useChat.ts handleMessages() (line 625), add a new if (streamMode === "cost_update") branch that updates cost state.

Step 5: Expose cost state in ChatContextType
Add threadCost (or similar) to ChatContextType (line 31-93) and expose it via ChatContext.

This is exactly the same path that metadata, values, messages, aborted, and mcp_sandbox_unreachable events follow.

4. Reusable UI Patterns

Pattern	Component	Location	Reuse for
Pill-button with popover	ThreadSandboxStatus	frontend/src/components/status/ThreadSandboxStatus.tsx	CostBadge pill in ChatUtilityRow
Expandable popover panel	Tasks popover in ChatUtilityRow	frontend/src/components/chat/ChatUtilityRow.tsx (lines 32-99)	CostBreakdownPanel
Settings card	SandboxSettings	frontend/src/components/settings/SandboxSettings.tsx	BudgetConfigSettings card on Settings page
Badge variants (outline/colored)	Badge (shadcn)	frontend/src/components/ui/badge.tsx	Budget threshold warning badges
Dialog	Dialog (shadcn)	frontend/src/components/ui/dialog.tsx	BudgetConfigDialog
Per-message metadata text	Message component	frontend/src/components/lists/ChatMessages.tsx (line 214-229)	Per-turn cost display
Badge with icon	SubagentBadge	frontend/src/components/badges/SubagentBadge.tsx	CostBadge styling reference
Sidebar thread metadata	ThreadItem	frontend/src/components/drawers/app-sidebar.tsx (line 280-289)	Cost column in thread list

5. Gap Analysis: What the Spec Accounts For vs. What Is Missing

Accounted for in the spec:

• Running cost badge in chat header (but not specific about ChatUtilityRow vs ChatNav)
• Cost breakdown panel (per-turn + per-phase)
• Budget configuration dialog
• Warning indicators at 75%/100%
• SSE cost_update event for real-time updates

Missing or underspecified frontend requirements:

1. Settings page for budget configuration: The spec mentions "Budget configuration dialog" but does not specify a persistent settings-page section. The Settings page (frontend/src/pages/settings/index.tsx) should get a BudgetSettings card (following the pattern of SandboxSettings, DefaultModelSettings). This is where users set default per-assistant or global budget caps.

2. Real-time cost during streaming: The spec says "cost badge updates within 1 second of each LLM response." The current streaming architecture sends messages events token-by-token but values events only at the end of a turn. The backend needs to emit cost_update events after each LLM call completes (not after each token), which aligns with how TurnMetrics is captured. The frontend just needs to handle this event type. However, the spec should clarify: does cost update mid-stream (after each sub-agent call / tool iteration) or only at the end of the full response? The multi-turn loop (tool_call -> tool_result -> next LLM call) means cost events should fire after each intermediate LLM call, not just the final one.

3. Historical cost view / analytics page: The spec mentions "historical cost data per thread" via API but has no dedicated frontend page/route for cross-thread analytics. This could be a follow-up, but a simple cost column in the sidebar ThreadItem (showing per-thread total cost) would be a minimal first step.

4. Thread cost on load: When loading an existing thread via useThread.loadThread() (line 75-136 in frontend/src/hooks/useThread.ts), there is no mechanism to fetch cost data. A separate API call to GET /v0/threads/{id}/metrics would be needed in loadThread() or as a lazy fetch triggered when the CostBadge component mounts.

5. Cost data shape in types: No frontend TypeScript types exist for cost data. New types needed:

   • TurnMetrics (turn_id, model, input_tokens, output_tokens, total_cost, duration_ms, phase)
   • ThreadCostSummary (total_cost, total_tokens, turn_count, per_phase, per_model)
   • RunBudget (max_cost_usd, max_tokens, action_on_exceed, is_active)
   • CostUpdateEvent (SSE event: cumulative_cost, turn_cost, budget_status)
     These should live in frontend/src/lib/entities/metrics.ts.

6. New API service file: A metricsService.ts is needed in frontend/src/lib/services/ for REST calls: getThreadMetrics(), getThreadTurns(), getBudgets(), createBudget(), deleteBudget().

7. Budget warning UX during streaming: The spec mentions pause action triggers HITL interrupt. On the frontend, this means the stream would emit a special event (budget-exceeded interrupt), and the UI needs to show a confirmation dialog ("Budget exceeded. Continue?") that calls the HITL approval endpoint. The current handleMessages() does not handle HITL interrupt events from the stream. This may require a new "budget_exceeded" or "hitl_interrupt" event type.

8. Cost formatting utility: A small utility function (formatCost(usdAmount: number): string -> "$0.0012" or "$1.23") with appropriate precision (4 decimal places for sub-cent, 2 for > $1) should be added to frontend/src/lib/utils/format.ts.

6. Recommended Implementation Order (Frontend)

1. Types & service layer: frontend/src/lib/entities/metrics.ts + frontend/src/lib/services/metricsService.ts
2. SSE plumbing: Extend stream.ts types, fetchStreamReader.ts parsers, useChat.ts handler
3. CostBadge component: frontend/src/components/badges/CostBadge.tsx -- pill in ChatUtilityRow
4. CostBreakdownPanel: Popover or Sheet triggered from CostBadge click
5. Per-message cost display: Extend Message component with per-turn cost from cost_update events
6. BudgetSettings: Settings page card for global/per-assistant budget configuration
7. Budget warning indicators: Color transitions on CostBadge (green/yellow/red)
8. Thread cost in sidebar: Small cost metadata in ThreadItem (sidebar thread list)
9. Tests: Vitest component tests for CostBadge, CostBreakdownPanel, BudgetSettings

7. Key File Paths for Implementation

Purpose	Path
SSE event types	frontend/src/lib/entities/stream.ts
SSE parser (both readers)	frontend/src/lib/utils/fetchStreamReader.ts
Stream source abstraction	frontend/src/lib/utils/streamSource.ts
Chat hook (event handling)	frontend/src/hooks/useChat.ts
Chat context provider	frontend/src/context/ChatContext.tsx
Chat utility row (badge host)	frontend/src/components/chat/ChatUtilityRow.tsx
Chat composer	frontend/src/components/chat/ChatComposer.tsx
Message component (per-turn cost)	frontend/src/components/lists/ChatMessages.tsx
Chat nav header	frontend/src/components/nav/ChatNav.tsx
Settings page	frontend/src/pages/settings/index.tsx
Sidebar thread list	frontend/src/components/drawers/app-sidebar.tsx
Existing badge component	frontend/src/components/ui/badge.tsx
Thread service (REST)	frontend/src/lib/services/threadService.ts
User settings service	frontend/src/lib/services/userSettingsService.ts
Format utilities	frontend/src/lib/utils/format.ts

[ryaneggz]

Backend Implementation Analysis

Comprehensive codebase review to identify reusable infrastructure, gaps in the current plan, and recommended adjustments.

---

1. Existing Token Capture Code (Reusable vs New)

Reusable — cache_metrics_middleware in backend/src/utils/middleware.py (L79-128)

This @wrap_model_call middleware already extracts exactly the data cost tracking needs:

• usage_metadata.input_tokens (L103)
• usage_metadata.output_tokens (L104)
• usage_metadata.total_tokens (L105)
• input_token_details.cache_creation and cache_read (L108-113)

It currently only logs this data. The new cost-tracking middleware should either:

• Option A (recommended): Extend cache_metrics_middleware to also persist TurnMetrics and emit the SSE event. This avoids a second middleware doing the same usage_metadata extraction.
• Option B: Create a separate cost_tracking_middleware that runs after cache_metrics_middleware and reads the same AIMessage.usage_metadata. This is cleaner separation but duplicates extraction logic.

Needs to be built:

• Wall-clock duration timing (time.perf_counter() before/after the await handler(request) call)
• Cost calculation from tokens + pricing lookup
• TurnMetrics persistence (async DB write)
• SSE cost_update event emission
• Phase detection logic (planning/execution/reflection/tool_call)
• The ContextSchema currently only holds model and user_id — it needs thread_id and assistant_id for metrics to associate with the right thread

2. Middleware Stack Placement

The middleware stack is assembled in init_default_middleware() at backend/src/utils/middleware.py (L322-346):

stack = [
    compaction_middleware,       # Summarizes long conversations
    add_ai_message_metadata,    # Adds model to AIMessage
    cache_metrics_middleware,    # Logs token counts
    retry_model,                # Retries on failure
    *pii_middleware(),           # PII detection
    AutoEvictMiddleware(...),   # Large tool result eviction
]

Cost accumulation should happen as a @wrap_model_call middleware positioned right after cache_metrics_middleware (or merged into it). This placement ensures:

• It runs after retry (so only successful calls are metered)
• It has access to the full ModelResponse with usage_metadata
• It runs before compaction (so compaction's own LLM call can also be metered if desired)

Important nuance: compaction_middleware at L178 also calls await handler(request) which triggers its own LLM call (using DEFAULT_COMPACTION_MODEL). If this call should be tracked, the cost middleware must wrap the handler inside compaction's middleware chain. The current stack ordering handles this correctly — compaction_middleware is first, so it wraps everything below it.

Budget enforcement should be a @wrap_model_call that runs before the actual handler, inserted near the top of the stack, so it can reject the call before tokens are consumed.

3. Pricing Data Structure

The ChatModels enum at backend/src/constants/llm.py (L15-68) uses a provider:model format (e.g., openai:gpt-5.2, anthropic:claude-sonnet-4). The pricing dictionary should use this exact same key format for direct lookup:

# backend/src/constants/pricing.py
MODEL_PRICING: dict[str, dict[str, float]] = {
    "openai:gpt-5.2": {"input_per_1m": 2.50, "output_per_1m": 10.00},
    "anthropic:claude-sonnet-4": {"input_per_1m": 3.00, "output_per_1m": 15.00},
    ...
}

Gap in the plan: The spec references input_per_1k but industry standard is per-1M tokens. Recommend using input_per_1m / output_per_1m to match provider documentation and reduce float precision issues.

Ollama/self-hosted models: get_ollama_models() at L70-86 dynamically discovers models with no fixed enum. Pricing for these should default to 0.00 (self-hosted = free) unless overridden in the DB.

Bedrock models: These have different pricing than direct API. The bedrock_converse: prefix makes them distinguishable, so pricing entries should include the full prefix string.

4. Persistence Strategy — LangGraph Store vs New Table

Current data persistence pattern: All domain entities (threads, assistants, memories, tools, projects, settings) use the LangGraph AsyncPostgresStore via BaseRepo (L13-71 in base_repo.py). Data is stored using store.aput(namespace=(user_id, entity_type), key=entity_id, value=dict). This is a key-value store with JSON values, indexed for semantic search.

For TurnMetrics — recommend a dedicated SQL table (Alembic migration), NOT the LangGraph store. Reasons:

• Aggregation queries: Thread cost summaries need SUM(total_cost) GROUP BY thread_id — this is a SQL primitive but impossible in the LangGraph store's KV API.
• Time-range queries: Historical cost analytics need WHERE created_at BETWEEN x AND y — again, native SQL.
• Volume: Metrics records accumulate per-LLM-call (potentially hundreds per thread). The LangGraph store's asearch() has a limit parameter and isn't designed for large-volume analytical queries.
• Indexing: Composite indexes on (thread_id, created_at), (user_id, created_at), (model) are essential for performance.

For RunBudget — the LangGraph store is actually fine here, using namespace (user_id, "budgets"). Budget configs are low-volume, per-user, and follow the same CRUD pattern as assistants/settings. This avoids a migration and reuses BaseRepo. However, a SQL table is also acceptable.

For PricingOverrides — use either approach, but a SQL table is simpler since overrides are system-wide (not per-user) and need exact key lookup by model string.

Database connection approach: The turn_metrics SQL table should be accessed via the existing AsyncSessionLocal from backend/src/services/db.py (L49), which provides an async SQLAlchemy session. This is the same engine used by the users and tokens tables (the only two existing SQL tables).

5. Budget Enforcement and Abort/Interrupt Integration

Existing abort mechanism (backend/src/services/abort.py):

• Uses Redis keys with 5-minute TTL
• Workers poll AbortService.check_abort_signal(thread_id) on each streaming chunk (L325 in tasks.py)
• Pattern: fire-and-forget signal via Redis, polling in the streaming loop

Existing HITL interrupt mechanism (backend/src/schemas/entities/hitl.py):

• Uses LangGraph's built-in interrupt() primitive
• Thread pauses at a checkpoint, frontend calls GET /threads/{id}/interrupts to discover pending interrupts
• User resumes via POST /threads/{id}/resume with a HumanDecision
• Interrupt types: accept, edit, response, reject

Recommended approach for budget enforcement:

For pause action — Do NOT use the abort mechanism. Instead, use the HITL interrupt pattern. Before each model call, the budget middleware should check cumulative cost. If exceeded, call interrupt() to pause the thread with an InterruptInfo containing:

• tool_name: "budget_check"
• tool_args: {"current_cost": X, "budget_limit": Y, "thread_id": Z}
• config.allowed_actions: [ACCEPT, REJECT] (accept = continue, reject = stop)

For stop action — Raise a BudgetExceededError (new exception) that the stream_generator and _execute_agent_stream catch and emit as an SSE error event, following the existing error emission pattern at L416-417 of stream.py.

For warn action — Emit a custom SSE event ("budget_warning", {...}) via the existing custom stream mode, without interrupting execution.

Gap in the plan: The spec says "cached running total, refreshed on each turn" for budget checks. This needs careful implementation:

• The _StreamState class in stream.py (L238-245) tracks files_map and todos_list during streaming. A similar running_cost field should be added.
• For distributed mode (run_agent_stream in tasks.py), the running cost must be tracked in the worker's _stream_chunks_to_redis loop (L291-368), not via Redis caching.
• For local streaming (stream_generator in stream.py), track in the _StreamState object.

6. Additional Gaps Identified

ContextSchema is too minimal (backend/src/schemas/contexts/__init__.py):

@dataclass
class ContextSchema:
    model: str
    user_id: Optional[str] = None

The cost-tracking middleware needs thread_id and assistant_id from the runtime context to associate metrics with the right thread. These are currently only in the RunnableConfig.configurable dict, not the context schema. Either:

• Expand ContextSchema to include thread_id and assistant_id
• Or access them from request.state or request.config["configurable"] in the middleware (the ModelRequest passed to @wrap_model_call has state and may have config access)

Two streaming paths need instrumentation:

1. Local streaming: stream_generator() in stream.py (L323-498)
2. Distributed streaming: _execute_agent_stream() in tasks.py (L371-718)

Both paths call construct_agent() and agent.astream(). Since cost tracking is a middleware on the agent itself (via init_default_middleware()), it will automatically instrument both paths. But the SSE event emission and _StreamState.running_cost tracking needs to be wired separately for each path.

Phase detection is non-trivial: The spec lists phases planning, execution, reflection, tool_call. The deepagents graph does not expose explicit phase labels. Approximation:

• tool_call — detectable when the AI message has tool_calls (L44 in middleware.py: not last_msg.tool_calls)
• execution — default for normal agent turns
• planning/reflection — would require inspecting the graph node name from the stream chunk (e.g., updates mode shows node names). For v1, recommend simplifying to just model_call and tool_call.

Real-time SSE emission: The handle_multi_mode function in stream.py already handles a custom stream mode (L229-230). The cost middleware should emit a ("custom", {"type": "cost_update", "data": {...}}) event. The frontend subscribes to custom events and filters by type. This requires adding "custom" to the DEFAULT_STREAM_MODES list or ensuring it's always included when cost tracking is active.

7. Summary of Recommended Approach

Component	Strategy	File Location
Token capture	Extend existing cache_metrics_middleware	backend/src/utils/middleware.py
Pricing data	Dict keyed by provider:model string, matching ChatModels enum values	backend/src/constants/pricing.py (new)
TurnMetrics storage	Dedicated SQL table via Alembic migration + SQLAlchemy model	backend/migrations/versions/0002_*.py (new)
RunBudget storage	LangGraph store namespace (user_id, "budgets") via BaseRepo	backend/src/repos/budget_repo.py (new)
Budget enforcement	@wrap_model_call middleware, uses HITL interrupt() for pause	backend/src/utils/metrics.py (new)
Cost SSE events	("custom", {"type": "cost_update", ...}) via existing stream infrastructure	Middleware emits, stream.py passes through
API endpoints	New route files under backend/src/routes/v0/	metrics.py, pricing.py, budgets.py (new)
ContextSchema	Add thread_id, assistant_id fields	backend/src/schemas/contexts/__init__.py

[ryaneggz]

SSE Streaming Design — Cost Data

Current Architecture Summary

Orchestra's SSE stream uses a tuple-based wire format: every event is ["mode", payload]. The backend emits events via handle_multi_mode() in backend/src/utils/stream.py, which dispatches on the first element of each chunk tuple. The frontend parses events in FetchStreamReader.parseEvent() / ResponseBodyReader.parseEvent() (in frontend/src/lib/utils/fetchStreamReader.ts), which switches on type and currently recognizes: metadata, messages, values, error, aborted. Unknown types return null and are silently dropped.

LangGraph already supports a "custom" stream mode — the backend's handle_multi_mode() passes ("custom", payload) tuples through unchanged (line 229-230 of stream.py). This is the designed extension point.

Recommended Approach for Cost Streaming

Use the "custom" stream mode with a sub-typed payload. No new stream mode is needed.

Backend — emit cost events:

# Inside middleware or callback after each LLM call:
from langgraph.types import Send
# Or via stream_writer if available in the node:
stream_writer(("custom", {
    "type": "cost_update",
    "data": {
        "turn_cost_usd": 0.0023,
        "cumulative_cost_usd": 0.0145,
        "input_tokens": 1200,
        "output_tokens": 340,
        "model": "anthropic:claude-sonnet-4-20250514",
        "phase": "execution",
    }
}))

This flows through handle_multi_mode() as-is (the "custom" case returns chunk untouched) and is serialized to data: ["custom", {"type": "cost_update", "data": {...}}]\n\n.

Frontend — extend the parser to recognize "custom" events:

1. Add "custom" to SSEEventType in frontend/src/lib/entities/stream.ts
2. Add a CustomEvent interface: { type: "custom"; data: { type: string; data: unknown } }
3. In parseEvent() (fetchStreamReader.ts), add a case "custom" that returns the event
4. In useChat.ts handleMessages(), add a handler for streamMode === "custom" that dispatches on payload[1].type:
   • "cost_update" → update a costState ref/state in the chat hook

Why not a new top-level stream mode?

• The "custom" mode is already wired end-to-end in the backend. Adding a new top-level mode (e.g., "cost") would require changes to VALID_STREAM_MODES, handle_multi_mode(), convert_messages(), and both parseEvent() implementations on the frontend. The "custom" sub-typing approach gives us a single extension point that all future features can share.

Cost event frequency: Emit once per LLM call completion (not per token). This keeps SSE bandwidth low while giving the frontend enough data for a running cost badge.

Distributed mode: The _stream_chunks_to_redis() function in backend/src/workers/tasks.py already writes all handle_multi_mode() output to Redis — custom events flow through without changes.

[ryaneggz]

Middleware Stack Ordering Analysis for Cost Tracking

Current Stack (from init_default_middleware)

The middleware list in init_default_middleware is passed to create_deep_agent, which appends it after deepagents' own internal middleware. The framework composes middleware so that first in list = outermost layer (wraps all subsequent middleware). The current Orchestra stack:

Position	Middleware	Type	Purpose
1	compaction_middleware	wrap_model_call	Summarize old messages before LLM call
2	add_ai_message_metadata	after_model	Attach model name to AI messages
3	cache_metrics_middleware	wrap_model_call	Log prompt-cache hit ratios
4	retry_model	wrap_model_call	Retry LLM calls up to 3 times
5-6	PIIMiddleware (x2)	Input filter	Block credit cards and API keys
7	AutoEvictMiddleware	awrap_tool_call	Evict large tool results to filesystem
(cond.)	AnthropicPromptCachingMiddleware	Various	Cache control headers (if non-default TTL)

Because these are appended via init_default_middleware(backend=backend) + middleware inside init_graph, they run after deepagents' internal stack (TodoList, Memory, Filesystem, SubAgent, Summarization, PromptCaching, PatchToolCalls).

Where Cost-Tracking Middleware Should Go

The cost-tracking middleware (metrics_capture_middleware from backend/src/utils/metrics.py) needs to:

1. Capture token counts, cost, and duration for every LLM call
2. Record metrics after retries settle (so you only record the successful call, not failed attempts)
3. Capture metrics on the post-compaction messages (so token counts reflect what was actually sent)

Recommended position: immediately after retry_model (position 5), before PII middleware.

stack = [
    compaction_middleware,        # 1. Compact messages first
    add_ai_message_metadata,     # 2. Tag AI messages
    cache_metrics_middleware,     # 3. Log cache hit ratios
    retry_model,                 # 4. Retry on failure
    metrics_capture_middleware,  # 5. NEW: Record tokens, cost, duration
    *pii_middleware(),           # 6-7. PII filtering
    AutoEvictMiddleware(...),    # 8. Evict large tool results
]

Rationale:

• After retry_model: Since retry_model wraps the handler and retries up to 3 times, placing cost tracking inside it (i.e., after it in the list) means cost tracking only captures the call that actually succeeds. If placed before retry_model, you'd capture all 3 attempts, which inflates cost reporting. Place it after so it only meters the attempt that the retry middleware decides to pass through.
• After cache_metrics_middleware: Cache metrics logs prompt-cache hit ratios. Cost tracking should be a separate concern. Keep them independent -- cache metrics is logging-only, cost tracking persists to DB and enforces budgets.
• Before PII: PII middleware modifies input content. Cost tracking should see the request as-is for accurate token counting. PII runs on input only, so the token count difference is negligible, but conceptually cost tracking is a metering concern that belongs above content filtering.

Budget Enforcement Consideration

Budget checks (pre-call: "has this thread exceeded its budget?") should run as a before_model hook or at the very top of a wrap_model_call, ideally as the outermost Orchestra middleware (position 1, before compaction). This way, if a thread is over budget, you short-circuit before spending tokens on compaction summaries. This could be a separate middleware or combined into the metrics middleware with a pre-call check.

Should cost tracking merge with cache_metrics_middleware?

Recommendation: Keep them separate for now, consider merging later.

• cache_metrics_middleware is pure logging (no DB writes, no side effects).
• Cost tracking persists records and enforces budgets -- fundamentally different concerns.
• However, they both read usage_metadata from the response. A future refactor could combine them into a single MetricsMiddleware class that handles logging, persistence, and budget enforcement in one pass. This avoids double-parsing of usage metadata.

[ryaneggz]

Migration Strategy Note

Schema Changes Required

The spec calls for three new SQL tables:

1. turn_metrics -- per-turn token/cost/duration records (turn_id, thread_id, assistant_id, user_id, model, input_tokens, output_tokens, total_tokens, input_cost, output_cost, total_cost, duration_ms, phase, created_at)
2. run_budgets -- budget configuration (thread_id or assistant_id, max_cost_usd, max_tokens, action_on_exceed, is_active)
3. pricing_overrides -- admin pricing overrides (model, input_per_1k, output_per_1k, currency)

Critical Overlap with Observability (#904)

• turn_metrics overlaps heavily with agent_spans from feat: Agent phase observability — structured traces with frontend timeline #904. Both store per-invocation token counts, cost, duration, model, and phase.
• Recommended approach: Do NOT create a separate turn_metrics table. Instead, let feat: Agent phase observability — structured traces with frontend timeline #904's agent_spans table serve as the source of truth for per-turn metrics. This migration should only create:
  • run_budgets table (budget enforcement)
  • pricing_overrides table (admin pricing config)
• Cost aggregation (ThreadCostSummary) should be a computed view or service-layer aggregation over agent_spans, not a separate table.
• If feat: Agent phase observability — structured traces with frontend timeline #904 is not yet landed when this starts, create turn_metrics as a standalone table, but design it to be mergeable with agent_spans later.

Migration File

• If landing after feat: Agent phase observability — structured traces with frontend timeline #904: backend/migrations/versions/0003_add_budgets_pricing.py, down_revision = "0002_add_agent_spans".
• If landing independently: backend/migrations/versions/0002_add_cost_tracking.py, down_revision = "0001".

Parallel Development Safety

• Moderate conflict risk with feat: Agent phase observability — structured traces with frontend timeline #904 due to overlapping data model. Coordinate on whether token/cost data lives in agent_spans or turn_metrics.
• run_budgets and pricing_overrides are unique to this spec and conflict with nothing.
• No changes to the Assistant Pydantic model needed (budget is per-thread or per-assistant but stored in its own table, not as an Assistant field).

Budget Enforcement Note

• The RunBudget.assistant_id field references assistants, but assistants live in LangGraph Store (not SQL). Use a plain text/UUID column with no SQL foreign key constraint. Enforce referential integrity at the application layer.

Recommendation

• Coordinate with feat: Agent phase observability — structured traces with frontend timeline #904 on the token/cost data model before implementing.
• Migration creates 2-3 new tables (depending on feat: Agent phase observability — structured traces with frontend timeline #904 coordination), all additive with no existing table modifications.
• Hardcoded pricing constants in backend/src/constants/pricing.py need no migration.

[ryaneggz]

Dependency Graph Validation

Position in Dependency Graph

Cost Tracking (#905) — INDEPENDENT (Layer 0)
  └─ soft dep from: Evaluator (#899) for budget cap integration
  └─ depended upon by: Agent Observability (#904) for cost_usd, Social Media Campaign (#902) for cost reporting

Declared Dependencies

• None declared — spec is self-contained.

Validated Dependencies

Correct. Cost Tracking introduces entirely new infrastructure:

• TurnMetrics, ThreadCostSummary, RunBudget schemas (new metrics.py)
• Pricing constants (new pricing.py)
• Metrics capture middleware (new utils/metrics.py)
• Budget enforcement service
• New database tables (turn_metrics, run_budgets, pricing_overrides)
• REST API endpoints at /v0/threads/{id}/metrics, /v0/pricing, /v0/budgets

None of these depend on Planner, Evaluator, or any other spec.

Implicit Dependencies Found

1. HITL integration — Budget enforcement with action_on_exceed: "pause" triggers an HITL interrupt. The existing HITL system (InterruptConfig, HumanDecision in hitl.py) is already shipped, so this is a satisfied implicit dependency. However, the current HITL system is tool-call oriented (tool_name, tool_args). A budget-exceeded interrupt is NOT a tool call — it's a system event. The HITL schema may need extension (e.g., a new interrupt type) or the budget interrupt may need to be shoehorned into the existing tool-based model. This is an unresolved interface design question.

2. LLMController modification — The metrics capture middleware wraps LLM invocations. This touches the same LLMController code path that Planner (feat: Planner Agent — auto-expand short prompts into full product specs #898) and Evaluator (feat: Evaluator Agent — GAN-inspired generator/evaluator feedback loop #899) modify. Merge conflicts likely if built in parallel.

3. SSE streaming — The spec mentions "real-time cost display via SSE streaming events." The existing streaming infrastructure (utils/stream.py, stream_generator) would need to emit cost events alongside message events. This is a modification to shared infrastructure that no other spec addresses.

Interface Contracts Needed

1. Phase enum — Cost Tracking uses phase: Literal["planning", "execution", "reflection", "tool_call"]. Agent Observability (feat: Agent phase observability — structured traces with frontend timeline #904) uses phase: Literal["planner", "generator", "evaluator", "handoff", "tool_call", "retrieval"]. These MUST be unified into a single shared enum. Current naming conflict:

   Cost Tracking	Observability	Should Be
   planning	planner	Pick one
   execution	generator	Pick one
   reflection	evaluator	Pick one
   tool_call	tool_call	Agreed
   (missing)	handoff	Add to both
   (missing)	retrieval	Add to both

2. Metrics-to-Span correlation — Cost Tracking captures TurnMetrics per LLM call. Observability captures AgentSpan per phase. These need a shared identifier (e.g., span_id on TurnMetrics or turn_id on AgentSpan) so the frontend can join cost data with trace data.

3. Budget interrupt type — How does a budget-exceeded event fit into the existing HITL InterruptInfo schema (which has tool_name and tool_args)? Options:

   • Use tool_name: "budget_exceeded" as a sentinel value
   • Extend InterruptInfo with an interrupt_type discriminator

   This should be decided before implementation.

Schema Conflict Risk

• High. Phase enum conflict with Observability. HITL schema extension question. SSE streaming modification touches shared infra. Multiple migrations from parallel specs.

Critical Path Assessment

On critical path for Social Media Campaign (#902) Week 3. Also provides the cost_usd field that Observability (#904) needs. Can be built in parallel with Planner/Evaluator (Layer 0) but migration ordering must be coordinated. Recommended: Wave 1.

[ryaneggz]

Testing Strategy Recommendations

Infrastructure Context

Cost tracking touches middleware (existing pattern in test_middleware_integration.py), database (migration + repo), API endpoints, SSE streaming events, and HITL interrupts (pattern in test_checkpoint_hitl.py). The conftest.py provides autouse mock for external APIs, so mocked LLM responses can include token usage metadata for testing.

Recommended Tests

Unit Tests (tests/unit/schemas/test_metrics.py)

• TurnMetrics: all field validations, phase enum values, cost fields are Decimal not float
• ThreadCostSummary: aggregation from list of TurnMetrics is correct
• RunBudget: at least one of thread_id or assistant_id required (custom validator), action_on_exceed enum validation

Unit Tests (tests/unit/utils/test_metrics_middleware.py) — CRITICAL

• Metrics middleware captures token counts from LLM response usage field
• Metrics middleware computes cost = tokens * pricing rate (test with known pricing values)
• Metrics middleware records duration_ms via time.perf_counter() (mock perf_counter to return known values)
• Metrics middleware fires-and-forgets DB write (verify async task is spawned, does not block handler)
• Metrics middleware handles missing usage field gracefully (some responses may lack token counts)

Unit Tests (tests/unit/constants/test_pricing.py)

• Pricing table has entries for all major models (OpenAI gpt-4o, gpt-4.1-mini, Anthropic claude-sonnet-4-20250514, etc.)
• Pricing lookup returns correct rates for known models
• Pricing lookup returns fallback/zero for unknown models
• Admin pricing overrides take precedence over hardcoded defaults

Unit Tests (tests/unit/services/test_budget_enforcement.py) — CRITICAL

• Budget check: cumulative cost under budget → allows execution
• Budget pause action: triggers HITL interrupt when cost exceeds max_cost_usd (use MockInterrupt pattern from test_checkpoint_hitl.py)
• Budget warn action: continues execution but emits warning event
• Budget stop action: halts execution, returns budget-exceeded error
• Budget check uses cached running total (not re-querying DB on every call)
• 75% threshold warning: emitted when cost reaches 75% of budget

Unit Tests (tests/unit/routes/test_metrics_routes.py)

• GET /threads/{id}/metrics: returns summary with correct aggregation
• GET /threads/{id}/metrics/turns: returns per-turn breakdown, ordered by created_at
• GET /pricing: lists all pricing entries
• POST /pricing: creates/updates pricing override
• DELETE /pricing/{model}: removes override, falls back to hardcoded

Integration Tests (tests/integration/test_cost_tracking.py)

• Full flow via async_client: POST /api/llm/stream → verify TurnMetrics record is created with correct model, token counts, and cost
• Budget enforcement flow: set budget → run agent → verify HITL interrupt at budget limit

Frontend Tests

• Cost badge component: renders $0.00 initially, updates on SSE cost events, color changes at thresholds
• Cost breakdown panel: renders per-turn table, per-phase summary
• Budget config form: validates max_cost_usd is positive number, action_on_exceed dropdown

How to Test Without Real LLM Calls

The existing conftest.py mocks include static usage fields in responses. For metrics testing, modify the mock response to include realistic token counts:

# In test-specific fixture or conftest override:
respx.post(url__regex=r"^https://api\.openai\.com/.*").mock(
    return_value=respx.MockResponse(status_code=200, json={
        # ... existing fields ...
        "usage": {"prompt_tokens": 150, "completion_tokens": 50, "total_tokens": 200},
    })
)

Gaps in Spec's Test Section

• Missing: Race condition test — concurrent LLM calls updating the same thread's running total
• Missing: Decimal precision test — cost calculations should not lose precision due to float arithmetic
• Missing: Test for cost tracking during evaluator loop — each iteration's cost should be independently tracked AND summed
• Missing: Test for SSE cost event format — verify the streaming event includes cost data the frontend can parse
• Missing: Test for pricing override CRUD with authorization — only admin users should be able to modify pricing

Shared Fixtures Needed

• sample_turn_metrics: factory creating TurnMetrics with configurable model/tokens/cost
• mock_pricing_table: a subset of pricing entries for test assertions
• budget_fixture: pre-configured RunBudget at various threshold levels