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!(
50        model.provider.as_str(),
51        "cloudflare-workers-ai" | "cloudflare-ai-gateway"
52    ) || matches_host(&model.base_url, CLOUDFLARE_API_HOST)
53        || matches_host(&model.base_url, CLOUDFLARE_AI_GATEWAY_HOST)
54}
55
56fn is_opencode(model: &Model) -> bool {
57    matches!(model.provider.as_str(), "opencode" | "opencode-go")
58        || matches_host(&model.base_url, OPENCODE_HOST)
59}
60
61/// The app-identifying headers for `model` (native `getDefaultAttributionHeaders`).
62pub fn default_attribution_headers(model: &Model) -> BTreeMap<String, String> {
63    let mut headers = BTreeMap::new();
64    if !enabled() {
65        return headers;
66    }
67    if is_openrouter(model) {
68        headers.insert("HTTP-Referer".into(), "https://pi.dev".into());
69        headers.insert("X-OpenRouter-Title".into(), "pi".into());
70        headers.insert("X-OpenRouter-Categories".into(), "cli-agent".into());
71    }
72    if is_nvidia_nim(model) {
73        headers.insert("X-BILLING-INVOKE-ORIGIN".into(), "Pi".into());
74    }
75    if is_cloudflare(model) {
76        headers.insert("User-Agent".into(), "rpi-coding-agent".into());
77    }
78    headers
79}
80
81/// opencode session headers (native `getSessionHeaders`).
82pub fn session_headers(model: &Model, session_id: Option<&str>) -> BTreeMap<String, String> {
83    let mut headers = BTreeMap::new();
84    let Some(session_id) = session_id.filter(|s| !s.is_empty()) else {
85        return headers;
86    };
87    if is_opencode(model) {
88        headers.insert("x-opencode-session".into(), session_id.to_string());
89        headers.insert("x-opencode-client".into(), "rpi".into());
90    }
91    headers
92}
93
94/// Merge session + attribution headers with caller-supplied sources (later
95/// wins), matching native `mergeProviderAttributionHeaders`.
96pub fn merge_provider_attribution_headers(
97    model: &Model,
98    session_id: Option<&str>,
99    extra: &[&BTreeMap<String, String>],
100) -> BTreeMap<String, String> {
101    let mut merged = session_headers(model, session_id);
102    for (k, v) in default_attribution_headers(model) {
103        merged.insert(k, v);
104    }
105    for source in extra {
106        for (k, v) in *source {
107            merged.insert(k.clone(), v.clone());
108        }
109    }
110    merged
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use rpi_ai::types::{Api, ModelCost};
117
118    fn model(provider: &str, base_url: &str) -> Model {
119        let mut m = Model::new("id", "id", Api::AnthropicMessages, provider, base_url);
120        m.cost = ModelCost::default();
121        m
122    }
123
124    #[test]
125    fn openrouter_gets_headers() {
126        let headers =
127            default_attribution_headers(&model("openrouter", "https://openrouter.ai/api/v1"));
128        assert_eq!(
129            headers.get("X-OpenRouter-Title").map(String::as_str),
130            Some("pi")
131        );
132    }
133
134    #[test]
135    fn openrouter_detected_by_host() {
136        let headers = default_attribution_headers(&model("custom", "https://openrouter.ai/api/v1"));
137        assert!(headers.contains_key("HTTP-Referer"));
138    }
139
140    #[test]
141    fn unrelated_provider_has_no_headers() {
142        let headers = default_attribution_headers(&model("anthropic", "https://api.anthropic.com"));
143        assert!(headers.is_empty());
144    }
145
146    #[test]
147    fn opencode_session_headers() {
148        let headers = session_headers(&model("opencode", "https://opencode.ai/x"), Some("s1"));
149        assert_eq!(
150            headers.get("x-opencode-session").map(String::as_str),
151            Some("s1")
152        );
153        let none = session_headers(&model("anthropic", "https://api.anthropic.com"), Some("s1"));
154        assert!(none.is_empty());
155    }
156
157    #[test]
158    fn host_parsing() {
159        assert_eq!(
160            host_of("https://user:pw@openrouter.ai:443/api").as_deref(),
161            Some("openrouter.ai")
162        );
163        assert_eq!(host_of("not a url"), None);
164    }
165}
rpi_cli/attribution.rs

rpi_cli/
attribution.rs
