返回 CodeWhale
local_ollama.rs
根目录 / crates / tui / src / local_ollama.rs
1 //! First-run / missing-key adoption of a live local Ollama catalog.
2 //!
3 //! Virgin sessions default to the DeepSeek costume (`deepseek-flash`). When a
4 //! real local daemon answers `GET /api/tags` (or the OpenAI-compat
5 //! `GET /v1/models` roster), the painted route must switch to a tag that
6 //! actually exists — never leave DeepSeek flash as the chrome while a live
7 //! local catalog is sitting on `:11434`.
8
9 use std::time::Duration;
10
11 use codewhale_config::catalog::{
12 CatalogOffering, CatalogSource, ProviderCatalogDelta, base_url_fingerprint, now_unix,
13 };
14 use serde::Deserialize;
15
16 use crate::config::{ApiProvider, Config, DEFAULT_OLLAMA_BASE_URL};
17
18 const TAGS_PROBE_TIMEOUT: Duration = Duration::from_secs(2);
19
20 /// Result of a successful local Ollama tags/models probe.
21 #[derive(Debug, Clone, PartialEq, Eq)]
22 pub(crate) struct LiveLocalOllamaCatalog {
23 pub(crate) endpoint_v1: String,
24 pub(crate) tags: Vec<String>,
25 }
26
27 impl LiveLocalOllamaCatalog {
28 /// Prefer the alphabetically first live tag (matches route_runtime's
29 /// Ollama default when tags have no `default_for_provider` flag).
30 pub(crate) fn preferred_tag(&self) -> Option<&str> {
31 self.tags.first().map(String::as_str)
32 }
33 }
34
35 /// True when this session should adopt a live local catalog into chrome.
36 ///
37 /// First-run and missing-key recovery paint DeepSeek by default; a live local
38 /// roster must replace that costume. An already-keyed hosted route is left alone.
39 #[must_use]
40 pub(crate) fn should_adopt_live_local_ollama(app: &crate::tui::app::App) -> bool {
41 if app.api_provider == ApiProvider::Ollama {
42 // Already on Ollama — route_runtime + #5795 own the tag; don't fight it.
43 return false;
44 }
45 app.onboarding_needs_api_key || app.onboarding_missing_key_recovery
46 }
47
48 /// Resolve the OpenAI-compat Ollama base URL (`…/v1`) from config defaults.
49 pub(crate) fn ollama_v1_base_url(config: &Config) -> String {
50 config
51 .provider_config_for(ApiProvider::Ollama)
52 .and_then(|entry| entry.base_url.clone())
53 .filter(|url| !url.trim().is_empty())
54 .unwrap_or_else(|| DEFAULT_OLLAMA_BASE_URL.to_string())
55 }
56
57 /// Strip a trailing `/v1` (with optional slash) so we can hit native `/api/tags`.
58 pub(crate) fn ollama_native_origin(v1_base: &str) -> String {
59 let trimmed = v1_base.trim().trim_end_matches('/');
60 if let Some(origin) = trimmed.strip_suffix("/v1") {
61 origin.to_string()
62 } else {
63 trimmed.to_string()
64 }
65 }
66
67 #[derive(Debug, Deserialize)]
68 struct OllamaTagsResponse {
69 #[serde(default)]
70 models: Vec<OllamaTagModel>,
71 }
72
73 #[derive(Debug, Deserialize)]
74 struct OllamaTagModel {
75 #[serde(default)]
76 name: String,
77 #[serde(default)]
78 model: String,
79 }
80
81 /// Parse native Ollama `GET /api/tags` JSON into sorted unique tag ids.
82 pub(crate) fn parse_ollama_tags_response(payload: &str) -> anyhow::Result<Vec<String>> {
83 let parsed: OllamaTagsResponse = serde_json::from_str(payload)
84 .map_err(|err| anyhow::anyhow!("Failed to parse Ollama /api/tags JSON: {err}"))?;
85 let mut tags: Vec<String> = parsed
86 .models
87 .into_iter()
88 .filter_map(|row| {
89 let name = row.name.trim();
90 if !name.is_empty() {
91 return Some(name.to_string());
92 }
93 let model = row.model.trim();
94 if !model.is_empty() {
95 Some(model.to_string())
96 } else {
97 None
98 }
99 })
100 .collect();
101 tags.sort();
102 tags.dedup();
103 Ok(tags)
104 }
105
106 fn record_ollama_tags_into_lake(endpoint_v1: &str, tags: &[String]) {
107 if tags.is_empty() {
108 return;
109 }
110 let fingerprint = base_url_fingerprint(endpoint_v1);
111 let fetched_at = now_unix();
112 let offerings = tags
113 .iter()
114 .map(|tag| CatalogOffering {
115 provider: "ollama".into(),
116 wire_model_id: tag.clone(),
117 endpoint_key: "chat".into(),
118 source: CatalogSource::Live {
119 base_url_fingerprint: fingerprint.clone(),
120 fetched_at,
121 },
122 default_for_provider: false,
123 ..Default::default()
124 })
125 .collect();
126 let ticket = crate::provider_catalog_live::begin_refresh_for_identity(
127 ApiProvider::Ollama,
128 "ollama",
129 endpoint_v1,
130 );
131 let _ = crate::provider_catalog_live::record_success_if_current(
132 &ticket,
133 ProviderCatalogDelta {
134 provider: "ollama".into(),
135 base_url_fingerprint: fingerprint,
136 fetched_at,
137 offerings,
138 },
139 );
140 }
141
142 async fn fetch_text(url: &str) -> anyhow::Result<String> {
143 // The first-run probe can run before any provider client has installed
144 // the rustls crypto provider; the shared builder installs it (the bare
145 // `reqwest::Client::builder()` panics under `rustls-no-provider`).
146 let client = crate::tls::reqwest_client_builder()
147 .timeout(TAGS_PROBE_TIMEOUT)
148 .build()?;
149 let response = client.get(url).send().await?;
150 if !response.status().is_success() {
151 anyhow::bail!("HTTP {}", response.status());
152 }
153 Ok(response.text().await?)
154 }
155
156 /// Probe local Ollama for a live catalog. Prefers native `/api/tags`, falls
157 /// back to OpenAI-compat `/v1/models`. Returns `None` when nothing useful
158 /// answered — never invents a tag.
159 pub(crate) async fn probe_live_local_ollama_catalog(
160 config: &Config,
161 ) -> Option<LiveLocalOllamaCatalog> {
162 let endpoint_v1 = ollama_v1_base_url(config);
163 let origin = ollama_native_origin(&endpoint_v1);
164 let tags_url = format!("{origin}/api/tags");
165
166 let tags = match fetch_text(&tags_url).await {
167 Ok(body) => match parse_ollama_tags_response(&body) {
168 Ok(tags) if !tags.is_empty() => tags,
169 Ok(_) => return None,
170 Err(err) => {
171 tracing::debug!(
172 target: "local_ollama",
173 error = %err,
174 "GET /api/tags returned unusable body"
175 );
176 Vec::new()
177 }
178 },
179 Err(err) => {
180 tracing::debug!(
181 target: "local_ollama",
182 error = %err,
183 url = %tags_url,
184 "GET /api/tags probe failed"
185 );
186 Vec::new()
187 }
188 };
189
190 let tags = if tags.is_empty() {
191 // Fallback: OpenAI-compat roster (same tags, different shape).
192 let models_url = format!("{}/models", endpoint_v1.trim_end_matches('/'));
193 match fetch_text(&models_url).await {
194 Ok(body) => match crate::client::parse_models_response(&body) {
195 Ok(models) if !models.is_empty() => models.into_iter().map(|m| m.id).collect(),
196 _ => return None,
197 },
198 Err(_) => return None,
199 }
200 } else {
201 tags
202 };
203
204 record_ollama_tags_into_lake(&endpoint_v1, &tags);
205 Some(LiveLocalOllamaCatalog { endpoint_v1, tags })
206 }
207
208 /// Env opt-out for harnesses that must not see the developer's machine.
209 ///
210 /// `spawn_local_ollama_adoption_probe` is already inert under `cfg(test)`, but
211 /// the PTY suites spawn the real binary, so that guard never reaches them. A
212 /// developer running Ollama on :11434 therefore gets the launch screen replaced
213 /// by a "Provider switched: deepseek -> ollama" notice, and the PTY tests that
214 /// wait for launch text fail on their machine while CI stays green. Sealing the
215 /// HOME is not enough, because this leak arrives over the loopback network
216 /// rather than through the filesystem.
217 pub(crate) const DISABLE_LOCAL_OLLAMA_PROBE_ENV: &str = "CODEWHALE_DISABLE_LOCAL_OLLAMA_PROBE";
218
219 fn local_ollama_probe_disabled() -> bool {
220 std::env::var_os(DISABLE_LOCAL_OLLAMA_PROBE_ENV).is_some_and(|value| !value.is_empty())
221 }
222
223 /// Background probe used by the event loop (mirrors `spawn_startup_version_check`).
224 pub(crate) fn spawn_local_ollama_adoption_probe(
225 config: &Config,
226 should_probe: bool,
227 ) -> Option<tokio::task::JoinHandle<Option<LiveLocalOllamaCatalog>>> {
228 if !should_probe || local_ollama_probe_disabled() {
229 return None;
230 }
231 #[cfg(test)]
232 {
233 let _ = config;
234 None
235 }
236 #[cfg(not(test))]
237 {
238 let config = config.clone();
239 Some(tokio::spawn(async move {
240 probe_live_local_ollama_catalog(&config).await
241 }))
242 }
243 }
244
245 #[cfg(test)]
246 mod tests {
247 use super::*;
248 use crate::test_support::{EnvVarGuard, lock_test_env};
249
250 #[test]
251 fn parse_ollama_tags_response_reads_name_field() {
252 let body = r#"{"models":[{"name":"qwen2.5:0.5b","model":"qwen2.5:0.5b","size":0}]}"#;
253 let tags = parse_ollama_tags_response(body).expect("parse");
254 assert_eq!(tags, vec!["qwen2.5:0.5b".to_string()]);
255 }
256
257 #[test]
258 fn parse_ollama_tags_response_sorts_and_dedups() {
259 let body = r#"{"models":[
260 {"name":"zeta:tag"},
261 {"name":"alpha:tag"},
262 {"name":"alpha:tag"}
263 ]}"#;
264 let tags = parse_ollama_tags_response(body).expect("parse");
265 assert_eq!(tags, vec!["alpha:tag".to_string(), "zeta:tag".to_string()]);
266 }
267
268 #[test]
269 fn ollama_native_origin_strips_v1() {
270 assert_eq!(
271 ollama_native_origin("http://localhost:11434/v1"),
272 "http://localhost:11434"
273 );
274 assert_eq!(
275 ollama_native_origin("http://127.0.0.1:11434/v1/"),
276 "http://127.0.0.1:11434"
277 );
278 }
279
280 #[test]
281 fn preferred_tag_is_alphabetically_first_after_sort() {
282 let mut tags = vec!["zeta:tag".into(), "alpha:tag".into()];
283 tags.sort();
284 let catalog = LiveLocalOllamaCatalog {
285 endpoint_v1: "http://localhost:11434/v1".into(),
286 tags,
287 };
288 assert_eq!(catalog.preferred_tag(), Some("alpha:tag"));
289 }
290
291 #[tokio::test]
292 async fn probe_live_local_ollama_catalog_reads_api_tags() {
293 let _lock = lock_test_env();
294 let _live = crate::provider_lake::lock_live_snapshot();
295 let home = tempfile::tempdir().unwrap();
296 let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path());
297 crate::provider_catalog_live::reset_cache_for_test();
298 crate::provider_lake::clear_live_snapshot();
299
300 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
301 let addr = listener.local_addr().unwrap();
302 std::thread::spawn(move || {
303 let (mut stream, _) = listener.accept().unwrap();
304 use std::io::{Read, Write};
305 let mut buf = [0u8; 1024];
306 let _ = stream.read(&mut buf);
307 let body = br#"{"models":[{"name":"qwen2.5:0.5b","model":"qwen2.5:0.5b"}]}"#;
308 let header = format!(
309 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
310 body.len()
311 );
312 stream.write_all(header.as_bytes()).unwrap();
313 stream.write_all(body).unwrap();
314 });
315
316 let endpoint = format!("http://{addr}/v1");
317 let mut config = Config::default();
318 config.provider_config_for_mut(ApiProvider::Ollama).base_url = Some(endpoint.clone());
319
320 let catalog = probe_live_local_ollama_catalog(&config)
321 .await
322 .expect("tags probe should succeed");
323 assert_eq!(catalog.endpoint_v1, endpoint);
324 assert_eq!(catalog.tags, vec!["qwen2.5:0.5b".to_string()]);
325 assert_eq!(catalog.preferred_tag(), Some("qwen2.5:0.5b"));
326 }
327 }
328
328 lines RUST