lean_ctx/proxy/
usage_accounting.rs

1//! Measured-cost plumbing: request-side opt-in (#1179) and response-header
2//! extraction (#1189).
3//!
4//! OpenRouter only reports the billed charge (`usage.cost`, and
5//! `cost_details.upstream_inference_cost` for BYOK) when the request opts in
6//! with `"usage": {"include": true}`. This module decides *when* the proxy may
7//! inject that opt-in and performs the injection.
8//!
9//! The gate is deliberately narrow: `usage` is not an OpenAI Chat Completions
10//! parameter — api.openai.com rejects unknown top-level fields with a 400, and
11//! other OpenAI-compatible upstreams (Azure, Groq, vLLM…) are not guaranteed
12//! to tolerate it either. Injection therefore only happens when the *effective*
13//! upstream host (post-routing) is openrouter.ai.
14//!
15//! Gateways that report the bill out-of-band do it via response headers:
16//! LiteLLM sends `x-litellm-response-cost` (USD) on every proxied turn, and
17//! corporate gateways often expose an equivalent under their own name
18//! (`[proxy] cost_response_header`). `cost_from_headers` turns those into
19//! the same measured figure the OpenRouter body path produces.
20
21use axum::http::HeaderMap;
22use serde_json::Value;
23
24/// Response header LiteLLM proxies attach to every turn: the billed USD.
25const LITELLM_COST_HEADER: &str = "x-litellm-response-cost";
26
27/// True when `base_url` points at OpenRouter — the only upstream documented to
28/// accept (and reward) the `usage.include` opt-in.
29pub(super) fn upstream_is_openrouter(base_url: &str) -> bool {
30    let rest = base_url
31        .strip_prefix("https://")
32        .or_else(|| base_url.strip_prefix("http://"))
33        .unwrap_or(base_url);
34    let host_port = rest.split(['/', '?']).next().unwrap_or(rest);
35    let host = host_port.split(':').next().unwrap_or(host_port);
36    host.eq_ignore_ascii_case("openrouter.ai")
37        || host.to_ascii_lowercase().ends_with(".openrouter.ai")
38}
39
40/// Billed USD reported by the upstream gateway via response headers (#1189):
41/// LiteLLM's standard header first, then the operator-configured extra header
42/// (already lowercase). Unparseable, negative or non-finite values are
43/// ignored — a broken gateway header must never poison the ledger.
44pub(super) fn cost_from_headers(headers: &HeaderMap, extra_header: Option<&str>) -> Option<f64> {
45    let mut names = vec![LITELLM_COST_HEADER];
46    if let Some(h) = extra_header {
47        names.push(h);
48    }
49    for name in names {
50        let cost = headers
51            .get(name)
52            .and_then(|v| v.to_str().ok())
53            .and_then(|s| s.trim().parse::<f64>().ok())
54            .filter(|c| c.is_finite() && *c >= 0.0);
55        if cost.is_some() {
56            return cost;
57        }
58    }
59    None
60}
61
62/// Injects `"usage": {"include": true}` into an OpenAI-shaped Chat Completions
63/// body so OpenRouter's final usage payload carries the billed `cost`.
64///
65/// Respects the caller: an existing `usage.include` (even `false`) is never
66/// overwritten, and a non-object `usage` value is left untouched.
67pub(super) fn inject_usage_include(doc: &mut Value) {
68    let Some(obj) = doc.as_object_mut() else {
69        return;
70    };
71    let usage = obj
72        .entry("usage")
73        .or_insert_with(|| Value::Object(serde_json::Map::new()));
74    if let Some(usage_obj) = usage.as_object_mut() {
75        usage_obj.entry("include").or_insert(Value::Bool(true));
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn openrouter_hosts_are_recognized() {
85        for url in [
86            "https://openrouter.ai/api",
87            "https://openrouter.ai",
88            "http://openrouter.ai:443/api",
89            "https://gateway.openrouter.ai/api",
90            "https://OPENROUTER.AI/api",
91        ] {
92            assert!(
93                upstream_is_openrouter(url),
94                "{url} must count as OpenRouter"
95            );
96        }
97    }
98
99    #[test]
100    fn other_upstreams_are_not_openrouter() {
101        for url in [
102            "https://api.openai.com",
103            "https://my-resource.services.ai.azure.com",
104            "https://api.groq.com/openai",
105            "http://127.0.0.1:11434",
106            "https://evil-openrouter.ai.example.com",
107            "https://notopenrouter.ai",
108        ] {
109            assert!(
110                !upstream_is_openrouter(url),
111                "{url} must NOT count as OpenRouter"
112            );
113        }
114    }
115
116    #[test]
117    fn litellm_cost_header_is_measured() {
118        let mut h = HeaderMap::new();
119        h.insert("x-litellm-response-cost", "0.00042".parse().unwrap());
120        assert_eq!(cost_from_headers(&h, None), Some(0.00042));
121    }
122
123    #[test]
124    fn operator_header_is_recognized_and_junk_ignored() {
125        let mut h = HeaderMap::new();
126        h.insert("x-corp-billed-usd", "1.25".parse().unwrap());
127        assert_eq!(
128            cost_from_headers(&h, Some("x-corp-billed-usd")),
129            Some(1.25),
130            "configured gateway header must be read"
131        );
132        assert_eq!(
133            cost_from_headers(&h, None),
134            None,
135            "unconfigured extra header is not consulted"
136        );
137
138        let mut junk = HeaderMap::new();
139        junk.insert("x-litellm-response-cost", "not-a-number".parse().unwrap());
140        junk.insert("x-corp-billed-usd", "-4".parse().unwrap());
141        assert_eq!(
142            cost_from_headers(&junk, Some("x-corp-billed-usd")),
143            None,
144            "unparseable and negative figures never enter the ledger"
145        );
146    }
147
148    #[test]
149    fn litellm_header_beats_extra_header_order() {
150        let mut h = HeaderMap::new();
151        h.insert("x-litellm-response-cost", "0.10".parse().unwrap());
152        h.insert("x-corp-billed-usd", "9.99".parse().unwrap());
153        assert_eq!(
154            cost_from_headers(&h, Some("x-corp-billed-usd")),
155            Some(0.10),
156            "the standard header wins when both are present"
157        );
158    }
159
160    #[test]
161    fn injects_usage_include_when_absent() {
162        let mut doc = serde_json::json!({"model": "deepseek/deepseek-v4-flash", "messages": []});
163        inject_usage_include(&mut doc);
164        assert_eq!(doc["usage"]["include"], Value::Bool(true));
165    }
166
167    #[test]
168    fn existing_opt_out_is_respected() {
169        let mut doc = serde_json::json!({"model": "m", "usage": {"include": false}});
170        inject_usage_include(&mut doc);
171        assert_eq!(doc["usage"]["include"], Value::Bool(false));
172    }
173
174    #[test]
175    fn non_object_usage_is_left_untouched() {
176        let mut doc = serde_json::json!({"model": "m", "usage": true});
177        inject_usage_include(&mut doc);
178        assert_eq!(doc["usage"], Value::Bool(true));
179    }
180}
lean_ctx/proxy/usage_accounting.rs

lean_ctx/proxy/
usage_accounting.rs
