rpi_cli/
attribution.rs

1//! Provider attribution headers.
2//!
3//! Port of native Pi's `packages/coding-agent/src/core/provider-attribution.ts`.
4//! Some providers want to know which app is calling them (for their own
5//! dashboards/quotas). Native pi attaches a small set of app-identifying headers
6//! for OpenRouter, NVIDIA NIM, and Cloudflare, plus a session header for
7//! opencode. This module reproduces that mapping.
8
9use std::collections::BTreeMap;
10
11use rpi_ai::model::Model;
12
13const OPENROUTER_HOST: &str = "openrouter.ai";
14const NVIDIA_NIM_HOST: &str = "integrate.api.nvidia.com";
15const CLOUDFLARE_API_HOST: &str = "api.cloudflare.com";
16const CLOUDFLARE_AI_GATEWAY_HOST: &str = "gateway.ai.cloudflare.com";
17const OPENCODE_HOST: &str = "opencode.ai";
18
19/// Attribution is enabled unless `RPI_DISABLE_ATTRIBUTION=1` (native gates this
20/// on install telemetry).
21pub fn enabled() -> bool {
22    !matches!(
23        std::env::var("RPI_DISABLE_ATTRIBUTION").as_deref(),
24        Ok("1") | Ok("true")
25    )
26}
27
28fn host_of(base_url: &str) -> Option<String> {
29    let rest = base_url.split("://").nth(1)?;
30    let host = rest.split(['/', '?', '#']).next()?;
31    let host = host.rsplit('@').next()?;
32    let host = host.split(':').next()?;
33    (!host.is_empty()).then(|| host.to_ascii_lowercase())
34}
35
36fn matches_host(base_url: &str, expected: &str) -> bool {
37    host_of(base_url).as_deref() == Some(expected)
38}
39
40fn is_openrouter(model: &Model) -> bool {
41    model.provider.as_str() == "openrouter" || model.base_url.contains(OPENROUTER_HOST)
42}
43
44fn is_nvidia_nim(model: &Model) -> bool {
45    model.provider.as_str() == "nvidia" || matches_host(&model.base_url, NVIDIA_NIM_HOST)
46}
47
48fn is_cloudflare(model: &Model) -> bool {
49    matches!(model.provider.as_str(), "cloudflare-workers-ai" | "cloudflare-ai-gateway")
50        || matches_host(&model.base_url, CLOUDFLARE_API_HOST)
51        || matches_host(&model.base_url, CLOUDFLARE_AI_GATEWAY_HOST)
52}
53
54fn is_opencode(model: &Model) -> bool {
55    matches!(model.provider.as_str(), "opencode" | "opencode-go")
56        || matches_host(&model.base_url, OPENCODE_HOST)
57}
58
59/// The app-identifying headers for `model` (native `getDefaultAttributionHeaders`).
60pub fn default_attribution_headers(model: &Model) -> BTreeMap<String, String> {
61    let mut headers = BTreeMap::new();
62    if !enabled() {
63        return headers;
64    }
65    if is_openrouter(model) {
66        headers.insert("HTTP-Referer".into(), "https://pi.dev".into());
67        headers.insert("X-OpenRouter-Title".into(), "pi".into());
68        headers.insert("X-OpenRouter-Categories".into(), "cli-agent".into());
69    }
70    if is_nvidia_nim(model) {
71        headers.insert("X-BILLING-INVOKE-ORIGIN".into(), "Pi".into());
72    }
73    if is_cloudflare(model) {
74        headers.insert("User-Agent".into(), "rpi-coding-agent".into());
75    }
76    headers
77}
78
79/// opencode session headers (native `getSessionHeaders`).
80pub fn session_headers(model: &Model, session_id: Option<&str>) -> BTreeMap<String, String> {
81    let mut headers = BTreeMap::new();
82    let Some(session_id) = session_id.filter(|s| !s.is_empty()) else {
83        return headers;
84    };
85    if is_opencode(model) {
86        headers.insert("x-opencode-session".into(), session_id.to_string());
87        headers.insert("x-opencode-client".into(), "rpi".into());
88    }
89    headers
90}
91
92/// Merge session + attribution headers with caller-supplied sources (later
93/// wins), matching native `mergeProviderAttributionHeaders`.
94pub fn merge_provider_attribution_headers(
95    model: &Model,
96    session_id: Option<&str>,
97    extra: &[&BTreeMap<String, String>],
98) -> BTreeMap<String, String> {
99    let mut merged = session_headers(model, session_id);
100    for (k, v) in default_attribution_headers(model) {
101        merged.insert(k, v);
102    }
103    for source in extra {
104        for (k, v) in *source {
105            merged.insert(k.clone(), v.clone());
106        }
107    }
108    merged
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use rpi_ai::types::{Api, ModelCost};
115
116    fn model(provider: &str, base_url: &str) -> Model {
117        let mut m = Model::new("id", "id", Api::AnthropicMessages, provider, base_url);
118        m.cost = ModelCost::default();
119        m
120    }
121
122    #[test]
123    fn openrouter_gets_headers() {
124        let headers = default_attribution_headers(&model("openrouter", "https://openrouter.ai/api/v1"));
125        assert_eq!(headers.get("X-OpenRouter-Title").map(String::as_str), Some("pi"));
126    }
127
128    #[test]
129    fn openrouter_detected_by_host() {
130        let headers = default_attribution_headers(&model("custom", "https://openrouter.ai/api/v1"));
131        assert!(headers.contains_key("HTTP-Referer"));
132    }
133
134    #[test]
135    fn unrelated_provider_has_no_headers() {
136        let headers = default_attribution_headers(&model("anthropic", "https://api.anthropic.com"));
137        assert!(headers.is_empty());
138    }
139
140    #[test]
141    fn opencode_session_headers() {
142        let headers = session_headers(&model("opencode", "https://opencode.ai/x"), Some("s1"));
143        assert_eq!(headers.get("x-opencode-session").map(String::as_str), Some("s1"));
144        let none = session_headers(&model("anthropic", "https://api.anthropic.com"), Some("s1"));
145        assert!(none.is_empty());
146    }
147
148    #[test]
149    fn host_parsing() {
150        assert_eq!(host_of("https://user:pw@openrouter.ai:443/api").as_deref(), Some("openrouter.ai"));
151        assert_eq!(host_of("not a url"), None);
152    }
153}
rpi_cli/attribution.rs

rpi_cli/
attribution.rs
