fix(usage): use OpenRouter's billed cost instead of estimating locally

daiv/automation/agent/base.py

@@ -152,6 +152,12 @@ def get_model(*, model: str, thinking_level: ThinkingLevel | None = None, **kwar
         model_kwargs = BaseAgent.get_model_kwargs(
             model=model_name, model_provider=provider, thinking_level=thinking_level, **kwargs
         )
+        if provider == ModelProvider.OPENROUTER:
+            # Route through our subclass so OpenRouter's `cost` field survives streaming.
+            from automation.agent.openrouter import ChatOpenRouter
+
+            model_kwargs.pop("model_provider", None)
+            return ChatOpenRouter(**model_kwargs)
         return init_chat_model(**model_kwargs)
 
     @staticmethod
@@ -207,18 +213,25 @@ def get_model_kwargs(
             }
             _kwargs["openai_api_base"] = site_settings.openrouter_api_base
             _kwargs["openai_api_key"] = site_settings.openrouter_api_key.get_secret_value()
+            # Opt into OpenRouter's authoritative usage accounting (returns the actual
+            # billed `cost` per call, including provider-specific cache-write rates that
+            # genai_prices cannot reconstruct from token counts alone). `stream_usage`
+            # ensures the OpenAI client sets `stream_options.include_usage=true`, which
+            # OpenRouter needs to emit usage on the final stream chunk.
+            _kwargs["extra_body"] = {"usage": {"include": True}}
+            _kwargs["stream_usage"] = True
 
             if thinking_level:
                 if _kwargs["model"].startswith(CLAUDE_THINKING_MODELS):
                     max_tokens, thinking_tokens = BaseAgent._get_anthropic_thinking_tokens(
                         thinking_level=thinking_level, max_tokens=_kwargs.get("max_tokens", CLAUDE_MAX_TOKENS)
                     )
                     _kwargs["max_tokens"] = max_tokens
-                    _kwargs["extra_body"] = {"reasoning": {"max_tokens": thinking_tokens}}
+                    _kwargs["extra_body"]["reasoning"] = {"max_tokens": thinking_tokens}
                     # When using thinking the temperature need to be set to 1 for Anthropic models
                     _kwargs["temperature"] = 1
                 else:
-                    _kwargs["extra_body"] = {"reasoning": {"effort": thinking_level.value}}
+                    _kwargs["extra_body"]["reasoning"] = {"effort": thinking_level.value}
 
             elif _kwargs["model"].startswith("anthropic") and "max_tokens" not in _kwargs:
                 # Avoid rate limiting by setting a fair max_tokens value

daiv/automation/agent/openrouter.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from langchain_openai.chat_models import ChatOpenAI
+
+if TYPE_CHECKING:
+    from langchain_core.outputs import ChatGenerationChunk
+
+
+class ChatOpenRouter(ChatOpenAI):
+    """``ChatOpenAI`` subclass that preserves OpenRouter-specific usage fields.
+
+    OpenRouter returns a ``cost`` field (USD) inside ``usage`` when the request
+    sets ``usage: {include: true}``. ``langchain_openai`` strips this when
+    converting raw usage to ``UsageMetadata`` while streaming, so the cost
+    is lost before our callbacks see it. This subclass stashes the cost on
+    the chunk's ``response_metadata`` so it survives chunk merging and is
+    reachable from the standard ``on_llm_end`` callback.
+
+    For non-streaming responses, the cost is already present at
+    ``message.response_metadata["token_usage"]["cost"]`` via the default
+    ``llm_output`` propagation, so no extra handling is needed there.
+    """
+
+    def _convert_chunk_to_generation_chunk(
+        self, chunk: dict, default_chunk_class: type, base_generation_info: dict | None
+    ) -> ChatGenerationChunk | None:
+        gen_chunk = super()._convert_chunk_to_generation_chunk(chunk, default_chunk_class, base_generation_info)
+        if gen_chunk is None:
+            return None
+        cost = (chunk.get("usage") or {}).get("cost")
+        if cost is not None:
+            gen_chunk.message.response_metadata["openrouter_cost_usd"] = str(cost)
+        return gen_chunk

daiv/automation/agent/usage_tracking.py

@@ -2,38 +2,105 @@
 
 import dataclasses
 import logging
+import threading
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass, field
-from decimal import Decimal
-from typing import TYPE_CHECKING, Any
+from decimal import Decimal, InvalidOperation
+from typing import TYPE_CHECKING, Any, override
 
 from genai_prices import Usage, calc_price
 from langchain_core.callbacks.usage import UsageMetadataCallbackHandler
+from langchain_core.messages import AIMessage
+from langchain_core.outputs import ChatGeneration
 from langchain_core.tracers.context import register_configure_hook
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
+    from langchain_core.outputs import LLMResult
+
 logger = logging.getLogger("daiv.usage")
 
+
+class DaivUsageCallbackHandler(UsageMetadataCallbackHandler):
+    """``UsageMetadataCallbackHandler`` extension that also captures provider-reported cost.
+
+    For OpenRouter responses, ``ChatOpenRouter`` stashes the billed ``cost`` (USD) on
+    each message's ``response_metadata`` (under ``openrouter_cost_usd`` for streaming, and
+    under ``token_usage.cost`` for non-streaming). This handler aggregates that value per
+    model, alongside the standard token usage metadata. ``build_usage_summary`` prefers
+    these provider-reported costs over local ``genai_prices`` calculations.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.cost_by_model: dict[str, Decimal] = {}
+        self._cost_lock = threading.Lock()
+
+    @override
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        super().on_llm_end(response, **kwargs)
+        try:
+            generation = response.generations[0][0]
+        except IndexError:
+            return
+        if not isinstance(generation, ChatGeneration):
+            return
+        message = generation.message
+        if not isinstance(message, AIMessage):
+            return
+        model_name = message.response_metadata.get("model_name")
+        if not model_name:
+            return
+        cost = _extract_provider_cost(message.response_metadata)
+        if cost is None:
+            return
+        with self._cost_lock:
+            self.cost_by_model[model_name] = self.cost_by_model.get(model_name, Decimal("0")) + cost
+
+
+def _extract_provider_cost(response_metadata: dict[str, Any]) -> Decimal | None:
+    """Pull the provider-reported cost out of an AIMessage's response_metadata.
+
+    Looks at the streaming-friendly key first (``openrouter_cost_usd``), then falls
+    back to the non-streaming ``token_usage.cost`` field exposed by langchain_openai's
+    ``llm_output`` propagation. Returns None when no cost is present or when the value
+    cannot be coerced to ``Decimal``.
+    """
+    candidates: list[Any] = []
+    if "openrouter_cost_usd" in response_metadata:
+        candidates.append(response_metadata["openrouter_cost_usd"])
+    token_usage = response_metadata.get("token_usage")
+    if isinstance(token_usage, dict) and "cost" in token_usage:
+        candidates.append(token_usage["cost"])
+    for raw in candidates:
+        if raw is None:
+            continue
+        try:
+            return Decimal(str(raw))
+        except InvalidOperation, ValueError:
+            logger.warning("Could not parse provider cost value %r", raw)
+    return None
+
+
 # Registered once at import time. Upstream ``get_usage_metadata_callback`` creates a new
 # ContextVar and hook registration on every call, leaking into the module-level hook list.
-_usage_metadata_var: ContextVar[UsageMetadataCallbackHandler | None] = ContextVar(
+_usage_metadata_var: ContextVar[DaivUsageCallbackHandler | None] = ContextVar(
     "daiv_usage_metadata_callback", default=None
 )
 register_configure_hook(_usage_metadata_var, inheritable=True)
 
 
 @contextmanager
-def track_usage_metadata() -> Iterator[UsageMetadataCallbackHandler]:
-    """Activate a ``UsageMetadataCallbackHandler`` for the enclosed block.
+def track_usage_metadata() -> Iterator[DaivUsageCallbackHandler]:
+    """Activate a ``DaivUsageCallbackHandler`` for the enclosed block.
 
     The handler is auto-propagated to every nested ``Runnable`` invocation (including
     subagents) via the registered ``ContextVar`` hook, so callers don't need to thread
     callbacks through ``RunnableConfig``.
     """
-    handler = UsageMetadataCallbackHandler()
+    handler = DaivUsageCallbackHandler()
     token = _usage_metadata_var.set(handler)
     try:
         yield handler
@@ -78,12 +145,21 @@ def _calc_model_cost(model_name: str, usage_metadata: dict[str, Any]) -> Decimal
     return result.total_price
 
 
-def build_usage_summary(handler_data: dict[str, dict[str, Any]]) -> UsageSummary:
-    """Build a UsageSummary from ``UsageMetadataCallbackHandler.usage_metadata``."""
+def build_usage_summary(
+    handler_data: dict[str, dict[str, Any]], provider_costs: dict[str, Decimal] | None = None
+) -> UsageSummary:
+    """Build a UsageSummary from ``UsageMetadataCallbackHandler.usage_metadata``.
+
+    When ``provider_costs`` contains an entry for a model, that authoritative cost is
+    used (e.g. OpenRouter's ``cost`` field) and the local ``genai_prices`` calculation
+    is skipped for that model. Models not in ``provider_costs`` fall back to
+    ``genai_prices``.
+    """
     if not handler_data:
         logger.warning("Usage metadata is empty; callback hook may not have fired on any LLM call")
         return UsageSummary()
 
+    provider_costs = provider_costs or {}
     total_input = 0
     total_output = 0
     total_total = 0
@@ -111,12 +187,18 @@ def build_usage_summary(handler_data: dict[str, dict[str, Any]]) -> UsageSummary
         if output_details := usage.get("output_token_details"):
             model_entry["output_token_details"] = dict(output_details)
 
-        model_cost = _calc_model_cost(model_name, usage)
-        if model_cost is not None:
-            model_entry["cost_usd"] = str(model_cost)
-            total_cost += model_cost
+        if (provider_cost := provider_costs.get(model_name)) is not None:
+            model_entry["cost_usd"] = str(provider_cost)
+            model_entry["cost_source"] = "provider"
+            total_cost += provider_cost
         else:
-            all_priced = False
+            model_cost = _calc_model_cost(model_name, usage)
+            if model_cost is not None:
+                model_entry["cost_usd"] = str(model_cost)
+                model_entry["cost_source"] = "genai_prices"
+                total_cost += model_cost
+            else:
+                all_priced = False
 
         by_model[model_name] = model_entry
 

daiv/codebase/managers/issue_addressor.py

@@ -144,7 +144,9 @@ async def _address_issue(self) -> AgentResult:
                     daiv_agent,
                     agent_config,
                     response=response_text,
-                    usage=build_usage_summary(usage_handler.usage_metadata).to_dict(),
+                    usage=build_usage_summary(
+                        usage_handler.usage_metadata, provider_costs=usage_handler.cost_by_model
+                    ).to_dict(),
                 )
 
     def _add_unable_to_address_issue_note(self, *, draft_published: bool = False):

daiv/codebase/managers/review_addressor.py

@@ -295,7 +295,9 @@ async def _address_comments(self) -> AgentResult:
                     daiv_agent,
                     agent_config,
                     response=response_text,
-                    usage=build_usage_summary(usage_handler.usage_metadata).to_dict(),
+                    usage=build_usage_summary(
+                        usage_handler.usage_metadata, provider_costs=usage_handler.cost_by_model
+                    ).to_dict(),
                 )
 
     def _add_unable_to_address_review_note(self, *, draft_published: bool = False):

daiv/jobs/tasks.py

@@ -61,5 +61,8 @@ async def run_job_task(repo_id: str, prompt: str, ref: str | None = None, use_ma
 
     logger.info("Job completed for repo_id=%s", repo_id)
     return await build_agent_result(
-        daiv_agent, config, response=response_text, usage=build_usage_summary(usage_handler.usage_metadata).to_dict()
+        daiv_agent,
+        config,
+        response=response_text,
+        usage=build_usage_summary(usage_handler.usage_metadata, provider_costs=usage_handler.cost_by_model).to_dict(),
     )

tests/unit_tests/automation/agent/test_openrouter.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from langchain_core.messages import AIMessageChunk
+
+from automation.agent.openrouter import ChatOpenRouter
+
+
+def _stream_chunk(*, content: str = "", usage: dict | None = None, finish_reason: str | None = None) -> dict:
+    """Build a raw OpenRouter SSE chunk dict matching the OpenAI streaming schema."""
+    chunk: dict = {
+        "id": "gen-1",
+        "model": "anthropic/claude-sonnet-4.6",
+        "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": finish_reason}],
+    }
+    if usage is not None:
+        chunk["usage"] = usage
+    return chunk
+
+
+class TestChatOpenRouterChunkConversion:
+    def _client(self) -> ChatOpenRouter:
+        # api_key required by the constructor; client is never actually used.
+        return ChatOpenRouter(
+            model="anthropic/claude-sonnet-4.6", api_key="sk-test", base_url="https://openrouter.test"
+        )
+
+    def test_stashes_cost_from_final_chunk(self):
+        client = self._client()
+        chunk = _stream_chunk(
+            usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150, "cost": 0.0987},
+            finish_reason="stop",
+        )
+
+        gen_chunk = client._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, base_generation_info=None)
+
+        assert gen_chunk is not None
+        assert gen_chunk.message.response_metadata["openrouter_cost_usd"] == "0.0987"
+
+    def test_no_cost_when_usage_missing(self):
+        """Mid-stream chunks have no usage block — leave response_metadata clean."""
+        client = self._client()
+        chunk = _stream_chunk(content="hello")
+
+        gen_chunk = client._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, base_generation_info=None)
+
+        assert gen_chunk is not None
+        assert "openrouter_cost_usd" not in gen_chunk.message.response_metadata
+
+    def test_no_cost_when_usage_lacks_cost_field(self):
+        """Older OpenRouter responses without `usage:include` won't have `cost`."""
+        client = self._client()
+        chunk = _stream_chunk(
+            usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, finish_reason="stop"
+        )
+
+        gen_chunk = client._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, base_generation_info=None)
+
+        assert gen_chunk is not None
+        assert "openrouter_cost_usd" not in gen_chunk.message.response_metadata