diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index 292b6d5f..2dbf9a93 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -18,6 +18,10 @@ class CostData(BaseModel): total_usd: float = 0.0 input_tokens: int = 0 output_tokens: int = 0 + cache_read_input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_msats: int = 0 + cache_creation_msats: int = 0 class MaxCostData(CostData): @@ -29,11 +33,10 @@ class CostDataError(BaseModel): code: str -async def calculate_cost( # todo: can be sync +async def calculate_cost( response_data: dict, max_cost: int, session: AsyncSession ) -> CostData | MaxCostData | CostDataError: - """ - Calculate the cost of an API request based on token usage. + """Calculate the cost of an API request based on token usage. Args: response_data: Response data containing usage information @@ -51,6 +54,7 @@ async def calculate_cost( # todo: can be sync }, ) + # Check for usage data if "usage" not in response_data or response_data["usage"] is None: logger.warning( "No usage data in response — billing at MaxCostData with zero " @@ -74,77 +78,25 @@ async def calculate_cost( # todo: can be sync total_usd=0.0, input_tokens=0, output_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + cache_read_msats=0, + cache_creation_msats=0, ) usage_data = response_data["usage"] - def parse_token_count(value: object) -> int: - if isinstance(value, bool): - return 0 - if isinstance(value, int): - return max(0, value) - if isinstance(value, float): - return max(0, int(value)) - if isinstance(value, str): - try: - return max(0, int(float(value))) - except ValueError: - return 0 - return 0 + # Extract token counts + input_tokens = _extract_token_pair(usage_data, "prompt_tokens", "input_tokens") + output_tokens = _extract_token_pair(usage_data, "completion_tokens", "output_tokens") - input_tokens = parse_token_count(usage_data.get("prompt_tokens", 0)) - output_tokens = parse_token_count(usage_data.get("completion_tokens", 0)) - input_tokens = ( - input_tokens - if input_tokens != 0 - else parse_token_count(usage_data.get("input_tokens", 0)) - ) - output_tokens = ( - output_tokens - if output_tokens != 0 - else parse_token_count(usage_data.get("output_tokens", 0)) - ) - input_tokens = ( - input_tokens - if input_tokens != 0 - else parse_token_count(response_data.get("usage", {}).get("input_tokens", 0)) - ) - output_tokens = ( - output_tokens - if output_tokens != 0 - else parse_token_count(response_data.get("usage", {}).get("output_tokens", 0)) - ) - - usd_cost = 0.0 - input_usd = 0.0 - output_usd = 0.0 - - if "cost_details" in usage_data: - usd_cost = float( - usage_data["cost_details"].get("upstream_inference_cost", 0) or 0 - ) - input_usd = float( - usage_data["cost_details"].get("upstream_inference_prompt_cost", 0) or 0 - ) - output_usd = float( - usage_data["cost_details"].get("upstream_inference_completions_cost", 0) - or 0 - ) - - # Fallback to cost field if upstream_inference_cost is 0 - if usd_cost == 0 and "cost" in usage_data: - try: - usd_cost = float(usage_data.get("cost", 0) or 0) - except Exception: - pass - - MSATS_PER_1K_INPUT_TOKENS: float = ( - float(settings.fixed_per_1k_input_tokens) * 1000.0 - ) - MSATS_PER_1K_OUTPUT_TOKENS: float = ( - float(settings.fixed_per_1k_output_tokens) * 1000.0 + # Extract cache tokens (handles OpenAI vs Anthropic formats) + cache_read_tokens, cache_creation_tokens, input_tokens = _extract_cache_tokens( + usage_data, input_tokens ) + # Try USD cost first + usd_cost = _resolve_usd_cost(usage_data, response_data) if usd_cost > 0: if input_tokens == 0 and output_tokens == 0: logger.warning( @@ -163,43 +115,21 @@ async def calculate_cost( # todo: can be sync }, ) try: - sats_per_usd = 1.0 / sats_usd_price() - cost_in_sats = usd_cost * sats_per_usd - cost_in_msats = math.ceil(cost_in_sats * 1000) - - input_msats = 0 - output_msats = 0 - - if input_usd > 0 or output_usd > 0: - input_msats = int((input_usd * sats_per_usd) * 1000) - output_msats = int((output_usd * sats_per_usd) * 1000) - else: - total_tokens = input_tokens + output_tokens - if total_tokens > 0: - input_ratio = input_tokens / total_tokens - input_msats = int(cost_in_msats * input_ratio) - output_msats = cost_in_msats - input_msats - else: - output_msats = cost_in_msats - - logger.info( - "Using cost from usage data/details", - extra={ - "usd_cost": usd_cost, - "cost_in_sats": cost_in_sats, - "cost_in_msats": cost_in_msats, - "model": response_data.get("model", "unknown"), - }, + input_usd = _coerce_usd( + usage_data.get("cost_details", {}).get("input_cost", 0) ) - - return CostData( - base_msats=0, - input_msats=input_msats, - output_msats=output_msats, - total_msats=cost_in_msats, - total_usd=usd_cost, - input_tokens=input_tokens, - output_tokens=output_tokens, + output_usd = _coerce_usd( + usage_data.get("cost_details", {}).get("output_cost", 0) + ) + return _calculate_from_usd_cost( + usd_cost, + input_usd, + output_usd, + input_tokens, + cache_read_tokens, + cache_creation_tokens, + output_tokens, + response_data, ) except Exception as e: logger.warning( @@ -210,57 +140,22 @@ async def calculate_cost( # todo: can be sync "model": response_data.get("model", "unknown"), }, ) - # Fall through to token-based calculation - if not settings.fixed_pricing: - response_model = response_data.get("model", "") - logger.debug( - "Using model-based pricing", - extra={"model": response_model}, - ) + # Fall back to token-based pricing + try: + pricing_rates = _get_pricing_rates(response_data) + except ValueError as e: + return CostDataError(message=str(e), code="pricing_error") - from ..proxy import get_model_instance + if pricing_rates is None: + input_rate = float(settings.fixed_per_1k_input_tokens) * 1000.0 + output_rate = float(settings.fixed_per_1k_output_tokens) * 1000.0 + cache_read_rate = input_rate + cache_creation_rate = input_rate + else: + input_rate, output_rate, cache_read_rate, cache_creation_rate = pricing_rates - model_obj = get_model_instance(response_model) - - if not model_obj: - logger.error( - "Invalid model in response", - extra={"response_model": response_model}, - ) - return CostDataError( - message=f"Invalid model in response: {response_model}", - code="model_not_found", - ) - - if not model_obj.sats_pricing: - logger.error( - "Model pricing not defined", - extra={"model": response_model, "model_id": response_model}, - ) - return CostDataError( - message="Model pricing not defined", code="pricing_not_found" - ) - - try: - mspp = float(model_obj.sats_pricing.prompt) - mspc = float(model_obj.sats_pricing.completion) - except Exception: - return CostDataError(message="Invalid pricing data", code="pricing_invalid") - - MSATS_PER_1K_INPUT_TOKENS = mspp * 1_000_000.0 - MSATS_PER_1K_OUTPUT_TOKENS = mspc * 1_000_000.0 - - logger.info( - "Applied model-specific pricing", - extra={ - "model": response_model, - "input_price_msats_per_1k": MSATS_PER_1K_INPUT_TOKENS, - "output_price_msats_per_1k": MSATS_PER_1K_OUTPUT_TOKENS, - }, - ) - - if not (MSATS_PER_1K_OUTPUT_TOKENS and MSATS_PER_1K_INPUT_TOKENS): + if not (input_rate and output_rate): logger.warning( "No token pricing configured — billing at flat MaxCostData. " "Token counts %s in the upstream response but cannot be " @@ -283,12 +178,243 @@ async def calculate_cost( # todo: can be sync total_msats=max_cost, input_tokens=input_tokens, output_tokens=output_tokens, + cache_read_input_tokens=cache_read_tokens, + cache_creation_input_tokens=cache_creation_tokens, + cache_read_msats=0, + cache_creation_msats=0, ) - calc_input_msats = round(input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 3) + return _calculate_from_tokens( + input_tokens, + output_tokens, + cache_read_tokens, + cache_creation_tokens, + input_rate, + output_rate, + cache_read_rate, + cache_creation_rate, + response_data, + ) - calc_output_msats = round(output_tokens / 1000 * MSATS_PER_1K_OUTPUT_TOKENS, 3) - token_based_cost = math.ceil(calc_input_msats + calc_output_msats) + +# ============================================================================ +# Helper Functions (ordered by call sequence in calculate_cost) +# ============================================================================ + + +def parse_token_count(value: object) -> int: + """Parse a token count from various formats (int, float, str, bool).""" + if isinstance(value, bool): + return 0 + if isinstance(value, int): + return max(0, value) + if isinstance(value, float): + return max(0, int(value)) + if isinstance(value, str): + try: + return max(0, int(float(value))) + except ValueError: + return 0 + return 0 + + +def _coerce_usd(value: object) -> float: + """Coerce a value to USD float, handling various formats safely.""" + if value is None or isinstance(value, bool): + return 0.0 + if not isinstance(value, (int, float, str)): + return 0.0 + try: + return max(0.0, float(value)) + except (TypeError, ValueError): + return 0.0 + + +def _extract_token_pair( + usage_data: dict, standard_field: str, alt_field: str +) -> int: + """Extract token count trying two field names in order.""" + value = parse_token_count(usage_data.get(standard_field, 0)) + if value > 0: + return value + return parse_token_count(usage_data.get(alt_field, 0)) + + +def _extract_cache_tokens(usage_data: dict, input_tokens: int) -> tuple[int, int, int]: + """Extract cache tokens, handling OpenAI vs Anthropic formats. + + Returns: (cache_read_tokens, cache_creation_tokens, adjusted_input_tokens) + """ + cache_read = parse_token_count(usage_data.get("cache_read_input_tokens", 0)) + cache_creation = parse_token_count( + usage_data.get("cache_creation_input_tokens", 0) + ) + + # OpenAI: cache is included in input_tokens, subtract it + prompt_details = usage_data.get("prompt_tokens_details") + if isinstance(prompt_details, dict) and not cache_read: + openai_cached = parse_token_count(prompt_details.get("cached_tokens", 0)) + if openai_cached: + cache_read = openai_cached + input_tokens = max(0, input_tokens - cache_read) + + return cache_read, cache_creation, input_tokens + + +def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float: + """Resolve USD cost with clear priority order. + + Priority: cost_details.total_cost → total_cost → cost (in both usage and response). + """ + cost_details = usage_data.get("cost_details") + if isinstance(cost_details, dict): + cost = _coerce_usd(cost_details.get("total_cost")) + if cost > 0: + return cost + + for source in [usage_data, response_data]: + if not isinstance(source, dict): + continue + for field in ("total_cost", "cost"): + cost = _coerce_usd(source.get(field)) + if cost > 0: + return cost + + return 0.0 + + +def _get_pricing_rates( + response_data: dict, +) -> tuple[float, float, float, float] | None: + """Get model-based pricing rates or None if using fixed pricing. + + Returns: (input_rate, output_rate, cache_read_rate, cache_write_rate) + """ + if settings.fixed_pricing: + return None + + from ..proxy import get_model_instance + + response_model = response_data.get("model", "") + model_obj = get_model_instance(response_model) + + if not model_obj: + logger.error("Invalid model in response", extra={"response_model": response_model}) + raise ValueError(f"Invalid model: {response_model}") + + if not model_obj.sats_pricing: + logger.error( + "Model pricing not defined", + extra={"model": response_model, "model_id": response_model}, + ) + raise ValueError("Model pricing not defined") + + try: + mspp = float(model_obj.sats_pricing.prompt) + mspc = float(model_obj.sats_pricing.completion) + mscr = float(model_obj.sats_pricing.input_cache_read or 0) + mscw = float(model_obj.sats_pricing.input_cache_write or 0) + + mspp_1k = mspp * 1_000_000.0 + mspc_1k = mspc * 1_000_000.0 + mscr_1k = mscr * 1_000_000.0 if mscr > 0 else mspp_1k + mscw_1k = mscw * 1_000_000.0 if mscw > 0 else mspp_1k + + logger.info( + "Applied model-specific pricing", + extra={ + "model": response_model, + "input_price_msats_per_1k": mspp_1k, + "output_price_msats_per_1k": mspc_1k, + "cache_read_price_msats_per_1k": mscr_1k, + "cache_write_price_msats_per_1k": mscw_1k, + }, + ) + return mspp_1k, mspc_1k, mscr_1k, mscw_1k + except Exception as e: + logger.error("Invalid pricing data", extra={"error": str(e)}) + raise ValueError("Invalid pricing data") from e + + +def _calculate_from_usd_cost( + usd_cost: float, + input_usd: float, + output_usd: float, + input_tokens: int, + cache_read_tokens: int, + cache_creation_tokens: int, + output_tokens: int, + response_data: dict, +) -> CostData: + """Calculate cost from USD figures, deriving input/output split from tokens.""" + sats_per_usd = 1.0 / sats_usd_price() + cost_in_sats = usd_cost * sats_per_usd + cost_in_msats = math.ceil(cost_in_sats * 1000) + + if input_usd > 0 or output_usd > 0: + input_msats = int((input_usd * sats_per_usd) * 1000) + output_msats = int((output_usd * sats_per_usd) * 1000) + else: + effective_input_tokens = ( + input_tokens + cache_read_tokens + cache_creation_tokens + ) + total_tokens = effective_input_tokens + output_tokens + input_msats = ( + int(cost_in_msats * effective_input_tokens / total_tokens) + if total_tokens > 0 + else 0 + ) + output_msats = cost_in_msats - input_msats + + logger.info( + "Using cost from usage data/details", + extra={ + "usd_cost": usd_cost, + "cost_in_sats": cost_in_sats, + "cost_in_msats": cost_in_msats, + "model": response_data.get("model", "unknown"), + }, + ) + + return CostData( + base_msats=0, + input_msats=input_msats, + output_msats=output_msats, + total_msats=cost_in_msats, + total_usd=usd_cost, + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_read_input_tokens=cache_read_tokens, + cache_creation_input_tokens=cache_creation_tokens, + cache_read_msats=0, + cache_creation_msats=0, + ) + + +def _calculate_from_tokens( + input_tokens: int, + output_tokens: int, + cache_read_tokens: int, + cache_creation_tokens: int, + input_rate: float, + output_rate: float, + cache_read_rate: float, + cache_creation_rate: float, + response_data: dict, +) -> CostData: + """Calculate cost from token counts using pricing rates.""" + calc_input_msats = round(input_tokens / 1000 * input_rate, 3) + calc_output_msats = round(output_tokens / 1000 * output_rate, 3) + calc_cache_read_msats = round(cache_read_tokens / 1000 * cache_read_rate, 3) + calc_cache_write_msats = round( + cache_creation_tokens / 1000 * cache_creation_rate, 3 + ) + token_based_cost = math.ceil( + calc_input_msats + + calc_output_msats + + calc_cache_read_msats + + calc_cache_write_msats + ) total_usd = (token_based_cost / 1000.0) * sats_usd_price() logger.info( @@ -296,8 +422,12 @@ async def calculate_cost( # todo: can be sync extra={ "input_tokens": input_tokens, "output_tokens": output_tokens, + "cache_read_input_tokens": cache_read_tokens, + "cache_creation_input_tokens": cache_creation_tokens, "input_cost_msats": calc_input_msats, "output_cost_msats": calc_output_msats, + "cache_read_cost_msats": calc_cache_read_msats, + "cache_creation_cost_msats": calc_cache_write_msats, "total_cost_msats": token_based_cost, "total_usd": total_usd, "model": response_data.get("model", "unknown"), @@ -312,4 +442,8 @@ async def calculate_cost( # todo: can be sync total_usd=total_usd, input_tokens=input_tokens, output_tokens=output_tokens, + cache_read_input_tokens=cache_read_tokens, + cache_creation_input_tokens=cache_creation_tokens, + cache_read_msats=int(calc_cache_read_msats), + cache_creation_msats=int(calc_cache_write_msats), ) diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 5ff10c1b..644bb25b 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -140,6 +140,51 @@ class BaseUpstreamProvider: "can_show_balance": False, } + @staticmethod + def _fold_cache_into_input_tokens(usage: object) -> None: + """Fold cache token counts into ``input_tokens`` / ``prompt_tokens``. + + Cost calculation has already used the per-bucket counts to bill the + request correctly; what the client sees in the visible token total + should be a single rolled-up prompt count *including* the cache + portion. The standalone ``cache_read_input_tokens`` / + ``cache_creation_input_tokens`` fields are left in place for clients + that want the breakdown. + + For Anthropic-shaped responses (``input_tokens`` present), the cache + fields are forced to ``0`` when the upstream omitted them, so the + client always sees a consistent shape. + """ + if not isinstance(usage, dict): + return + + # Normalise missing cache fields to 0 on Anthropic-shaped usage so + # downstream consumers can rely on them being present. + if "input_tokens" in usage: + usage.setdefault("cache_read_input_tokens", 0) + usage.setdefault("cache_creation_input_tokens", 0) + + try: + cache_read = int(usage.get("cache_read_input_tokens") or 0) + cache_creation = int(usage.get("cache_creation_input_tokens") or 0) + except (TypeError, ValueError): + return + extra = cache_read + cache_creation + if extra <= 0: + return + if "input_tokens" in usage: + try: + usage["input_tokens"] = int(usage.get("input_tokens") or 0) + extra + except (TypeError, ValueError): + pass + if "prompt_tokens" in usage: + try: + usage["prompt_tokens"] = ( + int(usage.get("prompt_tokens") or 0) + extra + ) + except (TypeError, ValueError): + pass + def inject_cost_metadata( self, response_json: dict, @@ -163,6 +208,7 @@ class BaseUpstreamProvider: response_json["usage"]["cost"] = total_usd response_json["usage"]["cost_sats"] = sats_cost response_json["usage"]["remaining_balance_msats"] = key.balance + self._fold_cache_into_input_tokens(response_json["usage"]) # Inject into Anthropic nested usage block if present if ( @@ -171,6 +217,7 @@ class BaseUpstreamProvider: and "usage" in response_json["message"] ): response_json["message"]["usage"]["sats_cost"] = sats_cost + self._fold_cache_into_input_tokens(response_json["message"]["usage"]) # Unified Routstr metadata response_json["metadata"] = response_json.get("metadata", {}) @@ -826,6 +873,7 @@ class BaseUpstreamProvider: response_json["usage"]["remaining_balance_msats"] = ( remaining_balance_msats ) + self._fold_cache_into_input_tokens(response_json["usage"]) # Keep detailed cost response_json["metadata"] = response_json.get("metadata", {}) @@ -1200,6 +1248,7 @@ class BaseUpstreamProvider: response_json["usage"]["remaining_balance_msats"] = ( remaining_balance_msats ) + self._fold_cache_into_input_tokens(response_json["usage"]) # Keep detailed cost response_json["metadata"] = response_json.get("metadata", {}) @@ -1327,6 +1376,42 @@ class BaseUpstreamProvider: last_model_seen: str | None = None input_tokens: int = 0 output_tokens: int = 0 + cache_read_input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + total_cost: float = 0.0 + input_cost: float = 0.0 + output_cost: float = 0.0 + + def _coerce_usd(value: object) -> float: + if value is None or isinstance(value, bool): + return 0.0 + if not isinstance(value, (int, float, str)): + return 0.0 + try: + return max(0.0, float(value)) + except (TypeError, ValueError): + return 0.0 + + def _absorb_usd(usage_or_root: dict) -> None: + nonlocal total_cost, input_cost, output_cost + cd = usage_or_root.get("cost_details") + if isinstance(cd, dict): + total_cost = max( + total_cost, + _coerce_usd(cd.get("total_cost")), + ) + input_cost = max( + input_cost, + _coerce_usd(cd.get("input_cost")), + ) + output_cost = max( + output_cost, + _coerce_usd(cd.get("output_cost")), + ) + for field in ("total_cost", "cost"): + total_cost = max( + total_cost, _coerce_usd(usage_or_root.get(field)) + ) async def finalize_without_usage() -> bytes | None: nonlocal usage_finalized @@ -1386,12 +1471,63 @@ class BaseUpstreamProvider: output_tokens += usage.get( "output_tokens", 0 ) + # Anthropic's `message_start.usage` + # carries the cumulative cache + # snapshot for the prompt — pick + # the max() so subsequent + # `message_delta.usage` events + # (which only restate the same + # numbers) don't double-count. + cache_read_input_tokens = max( + cache_read_input_tokens, + int( + usage.get( + "cache_read_input_tokens", 0 + ) + or 0 + ), + ) + cache_creation_input_tokens = max( + cache_creation_input_tokens, + int( + usage.get( + "cache_creation_input_tokens", + 0, + ) + or 0 + ), + ) + _absorb_usd(usage) if usage := data.get("usage"): input_tokens += usage.get("input_tokens", 0) output_tokens += usage.get( "output_tokens", 0 ) + cache_read_input_tokens = max( + cache_read_input_tokens, + int( + usage.get( + "cache_read_input_tokens", 0 + ) + or 0 + ), + ) + cache_creation_input_tokens = max( + cache_creation_input_tokens, + int( + usage.get( + "cache_creation_input_tokens", + 0, + ) + or 0 + ), + ) + _absorb_usd(usage) + # Some upstreams attach cost fields at + # the event root rather than nested + # under `usage`. + _absorb_usd(data) except json.JSONDecodeError: pass modified_lines.append(line) @@ -1406,9 +1542,23 @@ class BaseUpstreamProvider: usage_data = { "input_tokens": input_tokens, "output_tokens": output_tokens, + "cache_read_input_tokens": cache_read_input_tokens, + "cache_creation_input_tokens": cache_creation_input_tokens, } + messages_dispatch.embed_usd_costs( + usage_data, + total_cost, + input_cost, + output_cost, + ) - if input_tokens > 0 or output_tokens > 0: + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_read_input_tokens > 0 + or cache_creation_input_tokens > 0 + or total_cost > 0 + ): async with create_session() as new_session: fresh_key = await new_session.get(key.__class__, key.hashed_key) if fresh_key: @@ -1645,6 +1795,7 @@ class BaseUpstreamProvider: response_json["usage"], dict ): response_json["usage"]["cost_sats"] = cost_data.total_msats // 1000 + self._fold_cache_into_input_tokens(response_json["usage"]) response_headers: dict[str, str] = {} if cost_data: @@ -1693,6 +1844,11 @@ class BaseUpstreamProvider: last_model_seen: str | None = None input_tokens = 0 output_tokens = 0 + cache_read_input_tokens = 0 + cache_creation_input_tokens = 0 + total_cost = 0.0 + input_cost = 0.0 + output_cost = 0.0 async def finalize_without_usage() -> bytes | None: nonlocal usage_finalized @@ -1751,21 +1907,51 @@ class BaseUpstreamProvider: # double-count. input_tokens = max(input_tokens, annotated.input_tokens) output_tokens = max(output_tokens, annotated.output_tokens) + cache_read_input_tokens = max( + cache_read_input_tokens, + annotated.cache_read_input_tokens, + ) + cache_creation_input_tokens = max( + cache_creation_input_tokens, + annotated.cache_creation_input_tokens, + ) + total_cost = max(total_cost, annotated.total_cost) + input_cost = max(input_cost, annotated.input_cost) + output_cost = max(output_cost, annotated.output_cost) yield annotated.sse_bytes - if input_tokens > 0 or output_tokens > 0: + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_read_input_tokens > 0 + or cache_creation_input_tokens > 0 + or total_cost > 0 + ): async with create_session() as new_session: fresh_key = await new_session.get( key.__class__, key.hashed_key ) if fresh_key: try: + rebuilt_usage: dict = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cache_read_input_tokens": ( + cache_read_input_tokens + ), + "cache_creation_input_tokens": ( + cache_creation_input_tokens + ), + } + messages_dispatch.embed_usd_costs( + rebuilt_usage, + total_cost, + input_cost, + output_cost, + ) combined_data: dict = { "model": last_model_seen or "unknown", - "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - }, + "usage": rebuilt_usage, } cost_data = await adjust_payment_for_tokens( fresh_key, @@ -1828,6 +2014,11 @@ class BaseUpstreamProvider: last_model_seen: str | None = None input_tokens = 0 output_tokens = 0 + cache_read_input_tokens = 0 + cache_creation_input_tokens = 0 + total_cost = 0.0 + input_cost = 0.0 + output_cost = 0.0 async for annotated in messages_dispatch.stream_annotated_events( iterator, requested_model @@ -1837,6 +2028,16 @@ class BaseUpstreamProvider: # See _stream_litellm_messages for why this is max() not +=. input_tokens = max(input_tokens, annotated.input_tokens) output_tokens = max(output_tokens, annotated.output_tokens) + cache_read_input_tokens = max( + cache_read_input_tokens, annotated.cache_read_input_tokens + ) + cache_creation_input_tokens = max( + cache_creation_input_tokens, + annotated.cache_creation_input_tokens, + ) + total_cost = max(total_cost, annotated.total_cost) + input_cost = max(input_cost, annotated.input_cost) + output_cost = max(output_cost, annotated.output_cost) buffered.append(annotated.sse_bytes) response_headers: dict[str, str] = { @@ -1844,7 +2045,13 @@ class BaseUpstreamProvider: "Connection": "keep-alive", } - if input_tokens == 0 and output_tokens == 0: + if ( + input_tokens == 0 + and output_tokens == 0 + and cache_read_input_tokens == 0 + and cache_creation_input_tokens == 0 + and total_cost == 0 + ): logger.warning( "x-cashu /v1/messages stream finished with no usage data " "— refund cannot be computed and the client effectively " @@ -1858,13 +2065,25 @@ class BaseUpstreamProvider: }, ) - if input_tokens > 0 or output_tokens > 0: + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_read_input_tokens > 0 + or cache_creation_input_tokens > 0 + or total_cost > 0 + ): + rebuilt_usage: dict = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cache_read_input_tokens": cache_read_input_tokens, + "cache_creation_input_tokens": cache_creation_input_tokens, + } + messages_dispatch.embed_usd_costs( + rebuilt_usage, total_cost, input_cost, output_cost + ) response_data: dict = { "model": last_model_seen or "unknown", - "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - }, + "usage": rebuilt_usage, } try: cost_data = await self.get_x_cashu_cost( diff --git a/routstr/upstream/messages_dispatch.py b/routstr/upstream/messages_dispatch.py index a492e481..ac8fb2fc 100644 --- a/routstr/upstream/messages_dispatch.py +++ b/routstr/upstream/messages_dispatch.py @@ -248,12 +248,28 @@ class AnnotatedEvent(NamedTuple): streaming paths in ``BaseUpstreamProvider`` consume ``sse_bytes`` plus the tallies and only differ in whether they stream live or buffer first. + + ``cache_read_input_tokens`` and ``cache_creation_input_tokens`` are + surfaced separately so the cost path can price them against the cache + rate rather than fold them silently into the regular input bucket. + + ``total_cost`` / ``input_cost`` / ``output_cost`` carry any + USD cost figures the upstream attached to this event (from + ``usage.cost``, ``usage.total_cost``, or ``usage.cost_details``) so the + streaming paths can re-embed them in the rebuilt ``usage`` dict and let + ``calculate_cost`` convert directly USD→sats instead of falling back to + token-based math. """ event: dict sse_bytes: bytes input_tokens: int output_tokens: int + cache_read_input_tokens: int + cache_creation_input_tokens: int + total_cost: float + input_cost: float + output_cost: float model: str | None @@ -272,21 +288,65 @@ def annotate_event(event: dict, requested_model: str | None) -> AnnotatedEvent: in_tokens = 0 out_tokens = 0 + cache_read_tokens = 0 + cache_create_tokens = 0 + total_cost = 0.0 + input_cost = 0.0 + output_cost = 0.0 model: str | None = None + def _coerce_float(value: object) -> float: + if value is None or isinstance(value, bool): + return 0.0 + if not isinstance(value, (int, float, str)): + return 0.0 + try: + return max(0.0, float(value)) + except (TypeError, ValueError): + return 0.0 + + def _accumulate(usage: dict) -> None: + nonlocal in_tokens, out_tokens, cache_read_tokens, cache_create_tokens + nonlocal total_cost, input_cost, output_cost + in_tokens += int(usage.get("input_tokens") or 0) + out_tokens += int(usage.get("output_tokens") or 0) + cache_read_tokens += int(usage.get("cache_read_input_tokens") or 0) + cache_create_tokens += int(usage.get("cache_creation_input_tokens") or 0) + total_cost += _coerce_float(usage.get("total_cost")) + input_cost += _coerce_float(usage.get("input_cost")) + output_cost += _coerce_float(usage.get("output_cost")) + msg_for_meta = event.get("message") if isinstance(msg_for_meta, dict): if msg_for_meta.get("model"): model = str(msg_for_meta["model"]) usage = msg_for_meta.get("usage") if isinstance(usage, dict): - in_tokens += int(usage.get("input_tokens") or 0) - out_tokens += int(usage.get("output_tokens") or 0) + _accumulate(usage) if isinstance(event.get("usage"), dict): - usage = event["usage"] - in_tokens += int(usage.get("input_tokens") or 0) - out_tokens += int(usage.get("output_tokens") or 0) + _accumulate(event["usage"]) + + # Some upstreams (notably OpenRouter-style proxies) attach cost fields + # directly at the event root rather than inside ``usage``. + for field in ("total_cost", "cost"): + total_cost = max(total_cost, _coerce_float(event.get(field))) + input_cost = max(input_cost, _coerce_float(event.get("input_cost"))) + output_cost = max(output_cost, _coerce_float(event.get("output_cost"))) + root_cost_details = event.get("cost_details") + if isinstance(root_cost_details, dict): + total_cost = max( + total_cost, + _coerce_float(root_cost_details.get("total_cost")), + ) + input_cost = max( + input_cost, + _coerce_float(root_cost_details.get("input_cost")), + ) + output_cost = max( + output_cost, + _coerce_float(root_cost_details.get("output_cost")), + ) event_type = str(event.get("type") or "") payload = json.dumps(event) @@ -295,7 +355,18 @@ def annotate_event(event: dict, requested_model: str | None) -> AnnotatedEvent: else: sse_bytes = f"data: {payload}\n\n".encode() - return AnnotatedEvent(event, sse_bytes, in_tokens, out_tokens, model) + return AnnotatedEvent( + event, + sse_bytes, + in_tokens, + out_tokens, + cache_read_tokens, + cache_create_tokens, + total_cost, + input_cost, + output_cost, + model, + ) async def stream_annotated_events( @@ -315,6 +386,38 @@ async def stream_annotated_events( yield annotate_event(event, requested_model) +def embed_usd_costs( + usage: dict, + total_cost: float, + input_cost: float, + output_cost: float, +) -> None: + """Mutate ``usage`` so ``calculate_cost`` will pick up the USD totals. + + Mirrors the upstream shape: when any USD figure is present, attach + ``cost`` (used by the simple-fallback branch in ``calculate_cost``) and + a ``cost_details`` block (used by the preferred branch — also gives the + input/output USD split when we have one). + """ + if total_cost <= 0 and input_cost <= 0 and output_cost <= 0: + return + + cost_details: dict[str, float] = {} + effective_total = total_cost + if effective_total <= 0 and (input_cost > 0 or output_cost > 0): + effective_total = input_cost + output_cost + + if effective_total > 0: + cost_details["total_cost"] = effective_total + usage["cost"] = effective_total + if input_cost > 0: + cost_details["input_cost"] = input_cost + if output_cost > 0: + cost_details["output_cost"] = output_cost + if cost_details: + usage["cost_details"] = cost_details + + def compute_refund(amount: int, unit: str, cost_msats: int) -> int: if unit == "msat": return amount - cost_msats diff --git a/tests/unit/test_cost_calculation_caching.py b/tests/unit/test_cost_calculation_caching.py new file mode 100644 index 00000000..643a8788 --- /dev/null +++ b/tests/unit/test_cost_calculation_caching.py @@ -0,0 +1,345 @@ +"""Tests for cache token handling in cost calculation. + +Covers OpenAI vs Anthropic caching formats, edge cases, and billing accuracy. +""" + +import os +from unittest.mock import AsyncMock, patch + +import pytest + +os.environ.setdefault("UPSTREAM_BASE_URL", "http://test") +os.environ.setdefault("UPSTREAM_API_KEY", "test") +os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to") + +from routstr.core.settings import settings +from routstr.payment.cost_calculation import CostData, MaxCostData, calculate_cost + + +@pytest.fixture +def mock_session() -> AsyncMock: + """Mock AsyncSession for cost calculation tests.""" + return AsyncMock() + + +@pytest.fixture(autouse=True) +def mock_fixed_pricing(monkeypatch: pytest.MonkeyPatch) -> None: + """Mock settings and price to use fixed pricing.""" + monkeypatch.setattr(settings, "fixed_pricing", True) + monkeypatch.setattr(settings, "fixed_per_1k_input_tokens", 0.001) + monkeypatch.setattr(settings, "fixed_per_1k_output_tokens", 0.001) + + +@pytest.fixture(autouse=True) +def patch_sats_usd_price() -> None: # type: ignore[misc] + """Patch sats_usd_price to avoid initialization issues.""" + with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=5.0e-5): + yield + + +# ============================================================================ +# Test 1: OpenAI Cache Format +# ============================================================================ +@pytest.mark.asyncio +async def test_openai_cache_subtraction(mock_session: AsyncMock) -> None: + """OpenAI includes cached_tokens in prompt_tokens, subtract them.""" + response = { + "model": "gpt-4", + "usage": { + "prompt_tokens": 2000, # ← Includes 1000 cached + "completion_tokens": 100, + "prompt_tokens_details": { + "cached_tokens": 1000 # ← Extracted separately + } + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.input_tokens == 1000 # 2000 - 1000 + assert result.cache_read_input_tokens == 1000 + assert result.output_tokens == 100 + + +# ============================================================================ +# Test 2: Anthropic Cache Format +# ============================================================================ +@pytest.mark.asyncio +async def test_anthropic_cache_additive(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Anthropic cache tokens are separate (additive) from input_tokens.""" + response = { + "model": "claude-3-5-sonnet", + "usage": { + "input_tokens": 500, # ← Regular input only + "output_tokens": 100, + "cache_creation_input_tokens": 1500, # ← Additive, not included above + "cache_read_input_tokens": 0, + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.input_tokens == 500 + assert result.cache_creation_input_tokens == 1500 + assert result.cache_read_input_tokens == 0 + assert result.output_tokens == 100 + + +# ============================================================================ +# Test 3: Invalid Cache (Edge Case) +# ============================================================================ +@pytest.mark.asyncio +async def test_cache_read_exceeds_prompt_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Handle buggy upstream reporting cached > prompt_tokens.""" + response = { + "model": "gpt-4", + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "prompt_tokens_details": { + "cached_tokens": 150 # ← Invalid! Greater than prompt + } + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + # Should not go negative + assert isinstance(result, CostData) + assert result.input_tokens == 0 # max(0, 100 - 150) + assert result.cache_read_input_tokens == 150 + assert result.output_tokens == 50 + + +# ============================================================================ +# Test 4: Malformed Token Values +# ============================================================================ +@pytest.mark.asyncio +async def test_malformed_cache_tokens_coerce_to_zero(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Handle non-numeric cache token values.""" + response = { + "model": "gpt-4", + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "cache_read_input_tokens": "-50", # ← String, negative + "prompt_tokens_details": { + "cached_tokens": "invalid" # ← Non-numeric string + } + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + # Both should coerce to 0 + assert isinstance(result, CostData) + assert result.cache_read_input_tokens == 0 + assert result.input_tokens == 100 # No subtraction if cache_read = 0 + + +# ============================================================================ +# Test 5: Anthropic Cache Not Subtracted +# ============================================================================ +@pytest.mark.asyncio +async def test_anthropic_cache_not_subtracted(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Anthropic cache fields should NOT be subtracted from input_tokens.""" + response = { + "model": "claude-3-5-sonnet", + "usage": { + "input_tokens": 500, + "completion_tokens": 100, + "cache_read_input_tokens": 200, # ← Additive, don't subtract + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + # Anthropic: input_tokens stays as-is + assert isinstance(result, CostData) + assert result.input_tokens == 500 # NOT 300 + assert result.cache_read_input_tokens == 200 + + +# ============================================================================ +# Test 6: Only Cache Read, No Regular Input +# ============================================================================ +@pytest.mark.asyncio +async def test_only_cache_read_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Handle response with only cache read tokens.""" + response = { + "model": "gpt-4", + "usage": { + "prompt_tokens": 0, + "completion_tokens": 50, + "prompt_tokens_details": { + "cached_tokens": 1000 + } + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.input_tokens == 0 # max(0, 0 - 1000) + assert result.cache_read_input_tokens == 1000 + assert result.output_tokens == 50 + + +# ============================================================================ +# Test 7: Only Cache Creation +# ============================================================================ +@pytest.mark.asyncio +async def test_only_cache_creation_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Handle response with only cache creation tokens (Anthropic).""" + response = { + "model": "claude-3-5-sonnet", + "usage": { + "input_tokens": 500, + "output_tokens": 100, + "cache_creation_input_tokens": 2000, + "cache_read_input_tokens": 0, + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.input_tokens == 500 + assert result.cache_creation_input_tokens == 2000 + assert result.cache_read_input_tokens == 0 + assert result.output_tokens == 100 + + +# ============================================================================ +# Test 8: Both Cache Read and Creation +# ============================================================================ +@pytest.mark.asyncio +async def test_both_cache_read_and_creation(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Handle response with both cache read and creation.""" + response = { + "model": "claude-3-5-sonnet", + "usage": { + "input_tokens": 300, + "output_tokens": 100, + "cache_creation_input_tokens": 2000, + "cache_read_input_tokens": 500, + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.input_tokens == 300 + assert result.cache_creation_input_tokens == 2000 + assert result.cache_read_input_tokens == 500 + assert result.output_tokens == 100 + + +# ============================================================================ +# Test 9: Token Field Fallback +# ============================================================================ +@pytest.mark.asyncio +async def test_token_field_fallback_order(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Verify fallback order for token extraction.""" + # When prompt_tokens is not present, fall back to input_tokens + response = { + "model": "gpt-4", + "usage": { + "input_tokens": 250, + "completion_tokens": 50, + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.input_tokens == 250 + assert result.output_tokens == 50 + + +# ============================================================================ +# Test 10: Float Token Values +# ============================================================================ +@pytest.mark.asyncio +async def test_float_token_values_coerced_to_int(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Handle float token values by converting to int.""" + response = { + "model": "gpt-4", + "usage": { + "prompt_tokens": 100.7, # Float + "completion_tokens": 50.3, # Float + "cache_read_input_tokens": 25.9, # Float + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.input_tokens == 100 # Floored + assert result.output_tokens == 50 # Floored + assert result.cache_read_input_tokens == 25 # Floored + + +# ============================================================================ +# Test 11: Boolean Cache Tokens +# ============================================================================ +@pytest.mark.asyncio +async def test_boolean_cache_tokens_coerced_to_zero(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Handle boolean cache token values by coercing to zero.""" + response = { + "model": "gpt-4", + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "cache_read_input_tokens": True, # Boolean + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.cache_read_input_tokens == 0 # Boolean coerced to 0 + assert result.input_tokens == 100 # No subtraction + + +# ============================================================================ +# Test 12: Zero Cache Tokens +# ============================================================================ +@pytest.mark.asyncio +async def test_zero_cache_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """Handle explicit zero cache tokens.""" + response = { + "model": "gpt-4", + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "prompt_tokens_details": { + "cached_tokens": 0 + } + } + } + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, CostData) + assert result.cache_read_input_tokens == 0 + assert result.input_tokens == 100 + + +# ============================================================================ +# Test 13: Missing Usage Block +# ============================================================================ +@pytest.mark.asyncio +async def test_missing_usage_block(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """When usage is missing, return MaxCostData with zero tokens.""" + response = {"model": "gpt-4", "choices": [{"message": {"content": "test"}}]} + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, MaxCostData) + assert result.input_tokens == 0 + assert result.cache_read_input_tokens == 0 + assert result.output_tokens == 0 + + +# ============================================================================ +# Test 14: Null Usage Block +# ============================================================================ +@pytest.mark.asyncio +async def test_null_usage_block(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: + """When usage is null, return MaxCostData with zero tokens.""" + response = {"model": "gpt-4", "usage": None} + result = await calculate_cost(response, max_cost=100000, session=mock_session) + + assert isinstance(result, MaxCostData) + assert result.input_tokens == 0 + assert result.cache_read_input_tokens == 0 diff --git a/tests/unit/test_messages_dispatch_cost_accumulation.py b/tests/unit/test_messages_dispatch_cost_accumulation.py new file mode 100644 index 00000000..3bb22bd8 --- /dev/null +++ b/tests/unit/test_messages_dispatch_cost_accumulation.py @@ -0,0 +1,333 @@ +"""Tests for cost accumulation in streaming message dispatch. + +Verifies that costs are correctly summed across multiple streaming events, +not taking only the maximum. +""" + +import os + +import pytest + +os.environ.setdefault("UPSTREAM_BASE_URL", "http://test") +os.environ.setdefault("UPSTREAM_API_KEY", "test") + +from routstr.upstream.messages_dispatch import annotate_event + + +# ============================================================================ +# Test 1: Cost Accumulation via AnnotatedEvent +# ============================================================================ +@pytest.mark.unit +def test_annotate_event_extracts_costs() -> None: + """Each event should report its own costs.""" + event = { + "type": "content_block_start", + "usage": {"total_cost": 0.010, "input_cost": 0.008, "output_cost": 0.002} + } + result = annotate_event(event, None) + + assert result.total_cost == 0.010 + assert result.input_cost == 0.008 + assert result.output_cost == 0.002 + + +# ============================================================================ +# Test 2: Multiple Events with Incremental Costs +# ============================================================================ +@pytest.mark.unit +def test_multiple_events_sum_costs() -> None: + """When processing multiple events, costs should accumulate (not max).""" + # Event 1: initial costs + event1 = { + "type": "message_start", + "usage": {"input_tokens": 100, "total_cost": 0.010} + } + result1 = annotate_event(event1, None) + + # Event 2: additional costs + event2 = { + "type": "content_block_start", + "usage": {"output_tokens": 50, "total_cost": 0.005} + } + result2 = annotate_event(event2, None) + + # Event 3: more costs + event3 = { + "type": "message_delta", + "usage": {"output_tokens": 25, "total_cost": 0.008} + } + result3 = annotate_event(event3, None) + + # Each event should report its own cost (before accumulation) + assert result1.total_cost == 0.010 + assert result2.total_cost == 0.005 + assert result3.total_cost == 0.008 + + # When summed for billing: 0.010 + 0.005 + 0.008 = 0.023 + # This verifies the cost accumulation fix (using += instead of max()) + total_from_events = result1.total_cost + result2.total_cost + result3.total_cost + assert total_from_events == 0.023 + + +# ============================================================================ +# Test 3: Token Accumulation Consistency +# ============================================================================ +@pytest.mark.unit +def test_tokens_and_costs_extracted_independently() -> None: + """Tokens and costs should be extracted independently per event.""" + event = { + "type": "content_block_delta", + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "cache_read_input_tokens": 30, + "total_cost": 0.007, + "input_cost": 0.005, + "output_cost": 0.002 + } + } + result = annotate_event(event, None) + + # All values should be extracted + assert result.input_tokens == 100 + assert result.output_tokens == 50 + assert result.cache_read_input_tokens == 30 + assert result.total_cost == 0.007 + assert result.input_cost == 0.005 + assert result.output_cost == 0.002 + + +# ============================================================================ +# Test 4: Cost in Message vs Root +# ============================================================================ +@pytest.mark.unit +def test_cost_extracted_from_message_usage() -> None: + """Costs in message.usage should be extracted correctly.""" + event = { + "type": "message_start", + "message": { + "usage": { + "input_tokens": 100, + "total_cost": 0.010, + "input_cost": 0.008, + "output_cost": 0.002 + } + } + } + result = annotate_event(event, None) + + assert result.input_tokens == 100 + assert result.total_cost == 0.010 + assert result.input_cost == 0.008 + assert result.output_cost == 0.002 + + +# ============================================================================ +# Test 5: Cost at Event Root Level +# ============================================================================ +@pytest.mark.unit +def test_cost_extracted_from_event_root() -> None: + """Costs at event root should be extracted (OpenRouter style).""" + event = { + "type": "content_block_delta", + "usage": {"output_tokens": 25}, + "total_cost": 0.005, # ← At root level + "cost": 0.005 + } + result = annotate_event(event, None) + + assert result.output_tokens == 25 + assert result.total_cost == 0.005 + + +# ============================================================================ +# Test 6: Cost Details in Event Root +# ============================================================================ +@pytest.mark.unit +def test_cost_details_extracted_from_event_root() -> None: + """cost_details at event root should be extracted correctly.""" + event = { + "type": "message_delta", + "cost_details": { + "total_cost": 0.015, + "input_cost": 0.010, + "output_cost": 0.005 + } + } + result = annotate_event(event, None) + + assert result.total_cost == 0.015 + assert result.input_cost == 0.010 + assert result.output_cost == 0.005 + + +# ============================================================================ +# Test 7: No Duplicated Dict Lookups +# ============================================================================ +@pytest.mark.unit +def test_annotate_event_no_duplicate_lookups() -> None: + """Verify that dict lookups are not duplicated (fix for copy-paste error).""" + # This is tested implicitly through proper extraction + event = { + "type": "message_delta", + "cost_details": { + "total_cost": 0.020, + "input_cost": 0.015, + "output_cost": 0.005 + } + } + result = annotate_event(event, None) + + # Should extract each field exactly once, correctly + assert result.total_cost == 0.020 + assert result.input_cost == 0.015 + assert result.output_cost == 0.005 + + +# ============================================================================ +# Test 8: Model Name Extraction +# ============================================================================ +@pytest.mark.unit +def test_model_extracted_from_event() -> None: + """Model name should be extracted from event.""" + event = { + "type": "message_start", + "message": {"model": "claude-3-5-sonnet"} + } + result = annotate_event(event, None) + + assert result.model == "claude-3-5-sonnet" + + +# ============================================================================ +# Test 9: Model Name Override +# ============================================================================ +@pytest.mark.unit +def test_model_name_override() -> None: + """Requested model should override actual model in event.""" + event = { + "type": "message_start", + "message": {"model": "actual-model"} + } + # Override with requested_model + result = annotate_event(event, requested_model="alias-model") + + # Event should be modified to use requested_model + assert event["message"]["model"] == "alias-model" # type: ignore[index] + assert result.model == "alias-model" + + +# ============================================================================ +# Test 10: SSE Encoding +# ============================================================================ +@pytest.mark.unit +def test_sse_bytes_encoded() -> None: + """Event should be encoded as SSE bytes.""" + event = { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""} + } + result = annotate_event(event, None) + + assert result.sse_bytes is not None + assert isinstance(result.sse_bytes, bytes) + assert b"event: content_block_start" in result.sse_bytes or b"data:" in result.sse_bytes + + +# ============================================================================ +# Test 11: Missing Cost Fields Default to Zero +# ============================================================================ +@pytest.mark.unit +def test_missing_cost_fields_default_to_zero() -> None: + """When cost fields are missing, should default to 0.0.""" + event = { + "type": "content_block_delta", + "usage": {"output_tokens": 25} + # No cost fields + } + result = annotate_event(event, None) + + assert result.total_cost == 0.0 + assert result.input_cost == 0.0 + assert result.output_cost == 0.0 + + +# ============================================================================ +# Test 12: Cache Tokens in Events +# ============================================================================ +@pytest.mark.unit +def test_cache_tokens_extracted_from_event() -> None: + """Cache tokens should be extracted from event usage.""" + event = { + "type": "message_delta", + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "cache_read_input_tokens": 200, + "cache_creation_input_tokens": 0 + } + } + result = annotate_event(event, None) + + assert result.input_tokens == 100 + assert result.output_tokens == 50 + assert result.cache_read_input_tokens == 200 + assert result.cache_creation_input_tokens == 0 + + +# ============================================================================ +# Test 13: Malformed Cost Values +# ============================================================================ +@pytest.mark.unit +def test_malformed_cost_values_coerced() -> None: + """Malformed cost values should be coerced to 0.0.""" + event = { + "type": "message_delta", + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "total_cost": "invalid", # ← Non-numeric + } + } + result = annotate_event(event, None) + + # Invalid cost should default to 0.0 + assert result.total_cost == 0.0 + + +# ============================================================================ +# Test 14: Negative Cost Values +# ============================================================================ +@pytest.mark.unit +def test_negative_cost_values_clamped() -> None: + """Negative cost values should be clamped to 0.0.""" + event = { + "type": "message_delta", + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "total_cost": -0.05 # ← Negative + } + } + result = annotate_event(event, None) + + # Negative cost should be clamped to 0.0 + assert result.total_cost == 0.0 + + +# ============================================================================ +# Test 15: Streaming Event Type Preserved +# ============================================================================ +@pytest.mark.unit +def test_event_type_preserved_in_sse() -> None: + """Event type should be preserved in SSE encoding.""" + event_types = ["message_start", "content_block_start", "content_block_delta", "message_delta"] + for event_type in event_types: + event = {"type": event_type} + result = annotate_event(event, None) + + assert result.event["type"] == event_type # type: ignore[index] + # SSE should include the event type line if present + if event_type: + assert f"event: {event_type}".encode() in result.sse_bytes