1use std::collections::BTreeMap;
2
3use crate::backend::Ctx;
4use crate::model::ProviderId;
5
6pub fn configured_env(
7 ctx: &Ctx,
8 getenv: impl Fn(&str) -> Option<String>,
9) -> BTreeMap<String, String> {
10 let mut environment = BTreeMap::new();
11 for provider in [ProviderId::HuggingFace, ProviderId::ModelScope] {
12 let Some(config) = ctx.config.tool_sources(provider.as_str()) else {
13 continue;
14 };
15 if !config.env {
16 continue;
17 }
18 let sources = crate::model::source::effective_sources(ctx, provider);
19 let selected = config
20 .pin
21 .as_deref()
22 .and_then(|pin| sources.iter().find(|source| source.id == pin))
23 .or_else(|| sources.first());
24 let Some(source) = selected else {
25 continue;
26 };
27 let force = config.env_force;
28 match provider {
29 ProviderId::HuggingFace => {
30 let endpoint_managed = insert_if_allowed(
31 &mut environment,
32 "HF_ENDPOINT",
33 &source.download_url,
34 force,
35 &getenv,
36 );
37 let root = ctx
38 .dirs
39 .cache
40 .join("pkg")
41 .join("models")
42 .join("huggingface");
43 insert_path_if_allowed(&mut environment, "HF_HOME", &root, force, &getenv);
44 insert_path_if_allowed(
45 &mut environment,
46 "HF_HUB_CACHE",
47 &root.join("hub"),
48 force,
49 &getenv,
50 );
51 insert_path_if_allowed(
52 &mut environment,
53 "HF_XET_CACHE",
54 &root.join("xet"),
55 force,
56 &getenv,
57 );
58 insert_path_if_allowed(
59 &mut environment,
60 "HF_ASSETS_CACHE",
61 &root.join("assets"),
62 force,
63 &getenv,
64 );
65 if ctx.config.settings.offline {
66 insert_if_allowed(&mut environment, "HF_HUB_OFFLINE", "1", force, &getenv);
67 }
68 if endpoint_managed && !source.forward_credentials {
69 environment.insert("HF_HUB_DISABLE_IMPLICIT_TOKEN".into(), "1".into());
70 environment.insert("HF_TOKEN".into(), String::new());
71 environment.insert("HUGGING_FACE_HUB_TOKEN".into(), String::new());
72 environment.insert(
73 "HF_HOME".into(),
74 root.join("anonymous-home").display().to_string(),
75 );
76 }
77 }
78 ProviderId::ModelScope => {
79 let endpoint_managed = insert_if_allowed(
80 &mut environment,
81 "MODELSCOPE_ENDPOINT",
82 &source.download_url,
83 force,
84 &getenv,
85 );
86 let cache = ctx.dirs.cache.join("pkg").join("models").join("modelscope");
87 insert_path_if_allowed(
88 &mut environment,
89 "MODELSCOPE_CACHE",
90 &cache,
91 force,
92 &getenv,
93 );
94 if endpoint_managed && !source.forward_credentials {
95 environment.insert("MODELSCOPE_API_TOKEN".into(), String::new());
96 environment.insert(
97 "MODELSCOPE_HOME".into(),
98 cache.join("anonymous-home").display().to_string(),
99 );
100 }
101 }
102 }
103 }
104 environment
105}
106
107fn insert_path_if_allowed(
108 environment: &mut BTreeMap<String, String>,
109 key: &str,
110 value: &std::path::Path,
111 force: bool,
112 getenv: &impl Fn(&str) -> Option<String>,
113) -> bool {
114 insert_if_allowed(
115 environment,
116 key,
117 &value.display().to_string(),
118 force,
119 getenv,
120 )
121}
122
123fn insert_if_allowed(
124 environment: &mut BTreeMap<String, String>,
125 key: &str,
126 value: &str,
127 force: bool,
128 getenv: &impl Fn(&str) -> Option<String>,
129) -> bool {
130 if force || variable_is_available(key, getenv) {
131 environment.insert(key.to_string(), value.to_string());
132 true
133 } else {
134 false
135 }
136}
137
138fn variable_is_available(key: &str, getenv: &impl Fn(&str) -> Option<String>) -> bool {
139 let original_set = format!("OSDK_ORIG_{key}_SET");
140 if getenv(&original_set).is_some() {
141 return true;
142 }
143 getenv(key).is_none()
144}
145
146#[cfg(test)]
147mod tests {
148 use super::*;
149 use std::sync::Arc;
150
151 use crate::config::{Config, Settings, ToolSources};
152 use crate::dirs::Dirs;
153 use crate::platform::Platform;
154 use crate::source::Source;
155 use crate::store::Cas;
156
157 #[test]
158 fn global_huggingface_env_preserves_user_values_unless_forced() {
159 let temporary = tempfile::tempdir().unwrap();
160 let mut ctx = test_ctx(temporary.path());
161 ctx.config.sources.per_tool.insert(
162 "huggingface".into(),
163 ToolSources {
164 env: true,
165 ..Default::default()
166 },
167 );
168 let environment = configured_env(&ctx, |key| {
169 (key == "HF_ENDPOINT").then(|| "https://user.example".into())
170 });
171 assert!(!environment.contains_key("HF_ENDPOINT"));
172 assert!(environment.contains_key("HF_HUB_CACHE"));
173
174 ctx.config
175 .sources
176 .per_tool
177 .get_mut("huggingface")
178 .unwrap()
179 .env_force = true;
180 let environment = configured_env(&ctx, |key| {
181 (key == "HF_ENDPOINT").then(|| "https://user.example".into())
182 });
183 assert_eq!(
184 environment["HF_ENDPOINT"],
185 "https://huggingface.co".to_string()
186 );
187 }
188
189 #[test]
190 fn custom_endpoints_disable_implicit_credentials_and_restore_managed_values() {
191 let temporary = tempfile::tempdir().unwrap();
192 let mut ctx = test_ctx(temporary.path());
193 ctx.config.sources.per_tool.insert(
194 "huggingface".into(),
195 ToolSources {
196 pin: Some("custom".into()),
197 custom: vec![Source::mirror("custom", "https://mirror.example.test", 0)],
198 env: true,
199 ..Default::default()
200 },
201 );
202 let environment = configured_env(&ctx, |_| None);
203 assert_eq!(environment["HF_TOKEN"], "");
204 assert_eq!(environment["HF_HUB_DISABLE_IMPLICIT_TOKEN"], "1");
205 assert!(environment["HF_HOME"].contains("anonymous-home"));
206
207 let environment = configured_env(&ctx, |key| match key {
208 "HF_ENDPOINT" => Some("https://mirror.example.test".into()),
209 "OSDK_ORIG_HF_ENDPOINT_SET" => Some("1".into()),
210 "OSDK_ORIG_HF_ENDPOINT_PRESENT" => Some("1".into()),
211 "OSDK_ORIG_HF_ENDPOINT" => Some("https://user.example".into()),
212 _ => None,
213 });
214 assert_eq!(environment["HF_ENDPOINT"], "https://mirror.example.test");
215 }
216
217 #[test]
218 fn offline_is_only_exported_for_clients_that_support_it() {
219 let temporary = tempfile::tempdir().unwrap();
220 let mut ctx = test_ctx(temporary.path());
221 ctx.config.settings.offline = true;
222 for provider in ["huggingface", "modelscope"] {
223 ctx.config.sources.per_tool.insert(
224 provider.into(),
225 ToolSources {
226 env: true,
227 ..Default::default()
228 },
229 );
230 }
231 let environment = configured_env(&ctx, |_| None);
232 assert_eq!(environment["HF_HUB_OFFLINE"], "1");
233 assert!(!environment.contains_key("MODELSCOPE_OFFLINE"));
234 }
235
236 #[test]
237 fn custom_modelscope_endpoint_isolates_persisted_credentials() {
238 let temporary = tempfile::tempdir().unwrap();
239 let mut ctx = test_ctx(temporary.path());
240 ctx.config.sources.per_tool.insert(
241 "modelscope".into(),
242 ToolSources {
243 pin: Some("custom".into()),
244 custom: vec![Source::mirror(
245 "custom",
246 "https://modelscope.example.test",
247 0,
248 )],
249 env: true,
250 ..Default::default()
251 },
252 );
253 let environment = configured_env(&ctx, |key| {
254 (key == "MODELSCOPE_HOME").then(|| "/home/user/.modelscope".into())
255 });
256 assert_eq!(environment["MODELSCOPE_API_TOKEN"], "");
257 assert!(environment["MODELSCOPE_HOME"].contains("anonymous-home"));
258 assert_ne!(environment["MODELSCOPE_HOME"], "/home/user/.modelscope");
259 }
260
261 fn test_ctx(root: &std::path::Path) -> Ctx {
262 let dirs = Dirs::resolve_from(|key| match key {
263 "OSDK_DATA_DIR" => Some(root.join("data").display().to_string()),
264 "OSDK_CACHE_DIR" => Some(root.join("cache").display().to_string()),
265 "OSDK_CONFIG_DIR" => Some(root.join("config").display().to_string()),
266 _ => None,
267 })
268 .unwrap();
269 dirs.ensure().unwrap();
270 Ctx {
271 dirs: dirs.clone(),
272 platform: Platform::current(),
273 config: Config {
274 settings: Settings::default(),
275 sources: Default::default(),
276 tools: Default::default(),
277 tool_configs: Default::default(),
278 global_tools: Default::default(),
279 global_tool_configs: Default::default(),
280 tool_origins: Default::default(),
281 aliases: Default::default(),
282 project_config_path: None,
283 },
284 client: reqwest::Client::new(),
285 cas: Arc::new(Cas::new(dirs.store.clone())),
286 show_progress: false,
287 }
288 }
289}