1. 1//! Live model discovery: ask each provider what the supplied key can
2. 2//! actually reach, rather than shipping a curated table that drifts every
3. 3//! time a provider ships a model.
4. 4//!
5. 5//! Three response shapes cover every provider we speak to:
6. 6//! - OpenAI-compatible (`openai`, `openai-responses`, `openrouter`,
7. 7//! `openrouter-responses`, `opencode-zen`, `ollama`): `GET {base}/models`
8. 8//! → `{ "data": [{ "id" }] }`
9. 9//! - Anthropic: same path but `x-api-key` + `anthropic-version` headers.
10. 10//! - Google: `GET {base}/models?key=…` → `{ "models": [{ "name": "models/x" }] }`
11. 11
12. 12use std::time::Duration;
13. 13
14. 14use crate::error::LlmError;
15. 15use crate::registry::ProviderAuth;
16. 16
17. 17#[derive(Debug, Clone, PartialEq, Eq)]
18. 18pub struct ModelContext {
19. 19 pub input_tokens: u64,
20. 20 pub output_tokens: Option<u64>,
21. 21 /// Provider-reported quantisation label (Ollama `/api/show`
22. 22 /// `details.quantization_level`, e.g. "Q4_K_M"). `None` when the
23. 23 /// provider does not publish one — a `CapacityProfile` keys on this so a
24. 24 /// requantised model is measured fresh rather than inheriting a stale
25. 25 /// profile (docs/design/68-context-engine.md §1).
26. 26 pub quantisation: Option<String>,
27. 27}
28. 28
29. 29#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
30. 30pub struct BedrockModelAvailability {
31. 31 pub model_id: String,
32. 32 pub agreement_status: Option<String>,
33. 33 pub agreement_error: Option<String>,
34. 34 pub authorization_status: Option<String>,
35. 35 pub entitlement_status: Option<String>,
36. 36 pub region_status: Option<String>,
37. 37 pub invokable: bool,
38. 38}
39. 39
40. 40type BedrockAvailabilityCache =
41. 41 std::collections::HashMap<String, (std::time::Instant, Vec<BedrockModelAvailability>)>;
42. 42
43. 43/// Read Bedrock's native control-plane availability projection. This
44. 44/// intentionally remains separate from Mantle `/models`:
45. 45/// catalogue membership is not proof that a model can be invoked.
46. 46pub async fn bedrock_model_availability(
47. 47 auth: &ProviderAuth,
48. 48 model_ids: &[String],
49. 49) -> Result<Vec<BedrockModelAvailability>, LlmError> {
50. 50 use std::sync::{Mutex, OnceLock};
51. 51 static CACHE: OnceLock<Mutex<BedrockAvailabilityCache>> = OnceLock::new();
52. 52 let cache_key = format!(
53. 53 "{}:{:?}",
54. 54 auth.credential_id.as_deref().unwrap_or_default(),
55. 55 model_ids
56. 56 );
57. 57 if let Ok(cache) = CACHE
58. 58 .get_or_init(|| Mutex::new(std::collections::HashMap::new()))
59. 59 .lock()
60. 60 && let Some((at, value)) = cache.get(&cache_key)
61. 61 && at.elapsed() < Duration::from_secs(300)
62. 62 {
63. 63 return Ok(value.clone());
64. 64 }
65. 65 let region = auth
66. 66 .base_url
67. 67 .as_deref()
68. 68 .and_then(|url| url.split('.').nth(1))
69. 69 .unwrap_or("us-east-1");
70. 70 let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
71. 71 .region(aws_sdk_bedrock::config::Region::new(region.to_owned()))
72. 72 .load()
73. 73 .await;
74. 74 let client = aws_sdk_bedrock::Client::new(&config);
75. 75 let mut out = Vec::with_capacity(model_ids.len());
76. 76 for model_id in model_ids {
77. 77 let response = client
78. 78 .get_foundation_model_availability()
79. 79 .model_id(model_id)
80. 80 .send()
81. 81 .await
82. 82 .map_err(|e| LlmError::Api {
83. 83 status: 403,
84. 84 message: e.to_string(),
85. 85 })?;
86. 86 let agreement_status = response
87. 87 .agreement_availability()
88. 88 .map(|v| v.status().as_str().to_owned());
89. 89 let agreement_error = response
90. 90 .agreement_availability()
91. 91 .and_then(|v| v.error_message().map(str::to_owned));
92. 92 let authorization_status = Some(response.authorization_status().as_str().to_owned());
93. 93 let entitlement_status = Some(response.entitlement_availability().as_str().to_owned());
94. 94 let region_status = Some(response.region_availability().as_str().to_owned());
95. 95 out.push(BedrockModelAvailability {
96. 96 model_id: response.model_id().to_owned(),
97. 97 agreement_error,
98. 98 invokable: agreement_status.as_deref() == Some("AVAILABLE")
99. 99 && authorization_status.as_deref() == Some("AUTHORIZED")
100. 100 && entitlement_status.as_deref() == Some("AVAILABLE")
101. 101 && region_status.as_deref() == Some("AVAILABLE"),
102. 102 agreement_status,
103. 103 authorization_status,
104. 104 entitlement_status,
105. 105 region_status,
106. 106 });
107. 107 }
108. 108 if let Ok(mut cache) = CACHE
109. 109 .get_or_init(|| Mutex::new(std::collections::HashMap::new()))
110. 110 .lock()
111. 111 {
112. 112 cache.insert(cache_key, (std::time::Instant::now(), out.clone()));
113. 113 }
114. 114 Ok(out)
115. 115}
116. 116
117. 117const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(15);
118. 118/// Hard stop on paging so a malformed cursor can never loop forever.
119. 119const MAX_PAGES: usize = 20;
120. 120
121. 121fn http() -> Result<reqwest::Client, LlmError> {
122. 122 reqwest::Client::builder()
123. 123 .timeout(DISCOVERY_TIMEOUT)
124. 124 .build()
125. 125 .map_err(|e| LlmError::Network(e.to_string()))
126. 126}
127. 127
128. 128/// Map a non-success status onto the same error taxonomy the chat paths
129. 129/// use, so callers can distinguish "your key is wrong" from "provider is
130. 130/// down" without parsing strings.
131. 131fn status_error(status: u16, body: String) -> LlmError {
132. 132 let message = if body.trim().is_empty() {
133. 133 "no response body".to_string()
134. 134 } else {
135. 135 body.chars().take(400).collect()
136. 136 };
137. 137 match status {
138. 138 401 | 403 => LlmError::Auth(message),
139. 139 429 => LlmError::RateLimit {
140. 140 message,
141. 141 retry_after_secs: None,
142. 142 },
143. 143 400 => LlmError::classify_400(message),
144. 144 404 | 422 => LlmError::InvalidRequest(message),
145. 145 503 | 529 => LlmError::Overloaded(message),
146. 146 _ => LlmError::Api { status, message },
147. 147 }
148. 148}
149. 149
150. 150async fn read_json(res: reqwest::Response) -> Result<serde_json::Value, LlmError> {
151. 151 let status = res.status().as_u16();
152. 152 let body = res
153. 153 .text()
154. 154 .await
155. 155 .map_err(|e| LlmError::Network(e.to_string()))?;
156. 156 if !(200..300).contains(&status) {
157. 157 return Err(status_error(status, body));
158. 158 }
159. 159 serde_json::from_str(&body).map_err(|e| LlmError::Parse(e.to_string()))
160. 160}
161. 161
162. 162/// Ask `provider` which models `auth` unlocks. Returns ids sorted and
163. 163/// de-duplicated; never falls back to a baked-in list.
164. 164pub async fn list_models(provider: &str, auth: &ProviderAuth) -> Result<Vec<String>, LlmError> {
165. 165 let base = auth
166. 166 .base_url
167. 167 .as_deref()
168. 168 .filter(|b| !b.trim().is_empty())
169. 169 .or_else(|| default_base_url(provider))
170. 170 .ok_or_else(|| LlmError::InvalidRequest(format!("no base url for '{provider}'")))?
171. 171 .trim_end_matches('/')
172. 172 .to_string();
173. 173 let client = http()?;
174. 174
175. 175 let mut ids = match provider {
176. 176 // Anthropic pages with `has_more`/`last_id` and defaults to 20 per
177. 177 // page, so a single request would silently truncate the catalogue.
178. 178 "anthropic" => {
179. 179 // The chat base url carries no version segment, so add one here
180. 180 // rather than teaching every caller about it.
181. 181 let url = if base.ends_with("/v1") {
182. 182 format!("{base}/models")
183. 183 } else {
184. 184 format!("{base}/v1/models")
185. 185 };
186. 186 let mut out = Vec::new();
187. 187 let mut after: Option<String> = None;
188. 188 for _ in 0..MAX_PAGES {
189. 189 let mut req = client
190. 190 .get(&url)
191. 191 .header("x-api-key", &auth.api_key)
192. 192 .header("anthropic-version", crate::anthropic::ANTHROPIC_VERSION)
193. 193 .query(&[("limit", "1000")]);
194. 194 if let Some(cursor) = &after {
195. 195 req = req.query(&[("after_id", cursor.as_str())]);
196. 196 }
197. 197 let res = req
198. 198 .send()
199. 199 .await
200. 200 .map_err(|e| LlmError::Network(e.to_string()))?;
201. 201 let json = read_json(res).await?;
202. 202 out.extend(collect_data_ids(json.clone()));
203. 203 if json.get("has_more").and_then(|v| v.as_bool()) != Some(true) {
204. 204 break;
205. 205 }
206. 206 match json.get("last_id").and_then(|v| v.as_str()) {
207. 207 Some(cursor) => after = Some(cursor.to_string()),
208. 208 None => break,
209. 209 }
210. 210 }
211. 211 out
212. 212 }
213. 213 // Google pages with `nextPageToken` and defaults to 50 per page.
214. 214 "google" => {
215. 215 let mut out = Vec::new();
216. 216 let mut token: Option<String> = None;
217. 217 for _ in 0..MAX_PAGES {
218. 218 let mut req = client
219. 219 .get(format!("{base}/models"))
220. 220 .query(&[("key", auth.api_key.as_str()), ("pageSize", "1000")]);
221. 221 if let Some(cursor) = &token {
222. 222 req = req.query(&[("pageToken", cursor.as_str())]);
223. 223 }
224. 224 let res = req
225. 225 .send()
226. 226 .await
227. 227 .map_err(|e| LlmError::Network(e.to_string()))?;
228. 228 let json = read_json(res).await?;
229. 229 out.extend(collect_google_names(&json));
230. 230 match json.get("nextPageToken").and_then(|v| v.as_str()) {
231. 231 Some(cursor) if !cursor.is_empty() => token = Some(cursor.to_string()),
232. 232 _ => break,
233. 233 }
234. 234 }
235. 235 out
236. 236 }
237. 237 // Everything else speaks the OpenAI listing shape, which returns
238. 238 // the full set in one response.
239. 239 _ => {
240. 240 let res = client
241. 241 .get(format!("{base}/models"))
242. 242 .bearer_auth(&auth.api_key)
243. 243 .send()
244. 244 .await
245. 245 .map_err(|e| LlmError::Network(e.to_string()))?;
246. 246 collect_data_ids(read_json(res).await?)
247. 247 }
248. 248 };
249. 249
250. 250 ids.sort();
251. 251 ids.dedup();
252. 252 if ids.is_empty() {
253. 253 return Err(LlmError::Parse(
254. 254 "provider model catalogue contained no model ids".into(),
255. 255 ));
256. 256 }
257. 257 Ok(ids)
258. 258}
259. 259
260. 260fn context_from_json(provider: &str, json: &serde_json::Value) -> Option<ModelContext> {
261. 261 let data = json.get("data").unwrap_or(json);
262. 262 match provider {
263. 263 "google" => Some(ModelContext {
264. 264 input_tokens: data.get("inputTokenLimit")?.as_u64()?,
265. 265 output_tokens: data.get("outputTokenLimit").and_then(|v| v.as_u64()),
266. 266 quantisation: None,
267. 267 }),
268. 268 "ollama" => {
269. 269 // 1. Check modelfile parameters for "num_ctx <N>"
270. 270 let param_ctx = data
271. 271 .get("parameters")
272. 272 .and_then(|p| p.as_str())
273. 273 .and_then(|p| {
274. 274 for line in p.lines() {
275. 275 let mut parts = line.split_whitespace();
276. 276 if parts.next() == Some("num_ctx")
277. 277 && let Some(val) = parts.next().and_then(|v| v.parse::<u64>().ok())
278. 278 {
279. 279 return Some(val);
280. 280 }
281. 281 }
282. 282 None
283. 283 });
284. 284
285. 285 // 2. Check details.context_length
286. 286 let details_ctx = data
287. 287 .get("details")
288. 288 .and_then(|d| d.get("context_length"))
289. 289 .and_then(|v| v.as_u64());
290. 290
291. 291 // 3. Check model_info for any key ending with "context_length" (e.g. qwen35.context_length, gemma4.context_length)
292. 292 let model_info_ctx = data
293. 293 .get("model_info")
294. 294 .and_then(|mi| mi.as_object())
295. 295 .and_then(|obj| {
296. 296 obj.iter()
297. 297 .find(|(k, _)| k.ends_with(".context_length") || *k == "context_length")
298. 298 .and_then(|(_, v)| v.as_u64())
299. 299 });
300. 300
301. 301 let input_tokens = param_ctx.or(details_ctx).or(model_info_ctx).unwrap_or(8192);
302. 302
303. 303 // `/api/show` reports the running quantisation under
304. 304 // `details.quantization_level` (e.g. "Q4_K_M").
305. 305 let quantisation = data
306. 306 .get("details")
307. 307 .and_then(|d| d.get("quantization_level"))
308. 308 .and_then(|v| v.as_str())
309. 309 .map(str::to_string);
310. 310
311. 311 Some(ModelContext {
312. 312 input_tokens,
313. 313 output_tokens: Some(4096.min(input_tokens.saturating_div(2))),
314. 314 quantisation,
315. 315 })
316. 316 }
317. 317 _ => {
318. 318 let input_tokens = data
319. 319 .get("top_provider")
320. 320 .and_then(|v| v.get("context_length"))
321. 321 .and_then(|v| v.as_u64())
322. 322 .or_else(|| data.get("context_length").and_then(|v| v.as_u64()))?;
323. 323 Some(ModelContext {
324. 324 input_tokens,
325. 325 output_tokens: data
326. 326 .get("top_provider")
327. 327 .and_then(|v| v.get("max_completion_tokens"))
328. 328 .and_then(|v| v.as_u64()),
329. 329 quantisation: None,
330. 330 })
331. 331 }
332. 332 }
333. 333}
334. 334
335. 335/// Fetch the provider-reported context limits for one model. Providers that
336. 336/// do not publish machine-readable limits return `None`; callers must retain
337. 337/// their conservative configured limit in that case.
338. 338pub async fn model_context(
339. 339 provider: &str,
340. 340 auth: &ProviderAuth,
341. 341 model: &str,
342. 342) -> Result<Option<ModelContext>, LlmError> {
343. 343 let base = auth
344. 344 .base_url
345. 345 .as_deref()
346. 346 .filter(|b| !b.trim().is_empty())
347. 347 .or_else(|| default_base_url(provider))
348. 348 .ok_or_else(|| LlmError::InvalidRequest(format!("no base url for '{provider}'")))?
349. 349 .trim_end_matches('/');
350. 350 let client = http()?;
351. 351 let response = match provider {
352. 352 "google" => {
353. 353 client
354. 354 .get(format!("{base}/models/{model}"))
355. 355 .query(&[("key", auth.api_key.as_str())])
356. 356 .send()
357. 357 .await
358. 358 }
359. 359 "openrouter" => {
360. 360 client
361. 361 .get(format!("{base}/models/{model}"))
362. 362 .bearer_auth(&auth.api_key)
363. 363 .send()
364. 364 .await
365. 365 }
366. 366 "ollama" => {
367. 367 let root = base.trim_end_matches("/v1");
368. 368 let req = client
369. 369 .post(format!("{root}/api/show"))
370. 370 .json(&serde_json::json!({ "name": model }));
371. 371 let req = if !auth.api_key.is_empty() {
372. 372 req.bearer_auth(&auth.api_key)
373. 373 } else {
374. 374 req
375. 375 };
376. 376 match req.send().await {
377. 377 Ok(res) if res.status().is_success() => {
378. 378 let json = read_json(res).await?;
379. 379 return Ok(context_from_json("ollama", &json));
380. 380 }
381. 381 _ => {
382. 382 return Ok(Some(ModelContext {
383. 383 input_tokens: 8192,
384. 384 output_tokens: Some(4096),
385. 385 quantisation: None,
386. 386 }));
387. 387 }
388. 388 }
389. 389 }
390. 390 _ => return Ok(None),
391. 391 }
392. 392 .map_err(|e| LlmError::Network(e.to_string()))?;
393. 393 Ok(context_from_json(provider, &read_json(response).await?))
394. 394}
395. 395
396. 396/// Anthropic per-model capability flags this build conditions behaviour on
397. 397/// (docs/design/68-context-engine.md §11 "Anthropic" row). The Models API's
398. 398/// `capabilities` object carries many more fields; these are the ones an
399. 399/// adapter actually reads before it decides to send an optional parameter.
400. 400#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
401. 401pub struct AnthropicCapabilities {
402. 402 pub effort: bool,
403. 403 /// Best-effort key name — Anthropic's public docs describe the
404. 404 /// `capabilities` tree's `effort`/`thinking`/`image_input` members but
405. 405 /// do not (yet) enumerate a fast-mode entry. Absent or unrecognized
406. 406 /// shapes read as `false`, which only ever means "do not try fast mode
407. 407 /// for this model" — never a hard failure (invariant 9: unknown stays
408. 408 /// explicitly unavailable, it is never guessed into `true`).
409. 409 pub fast_mode: bool,
410. 410}
411. 411
412. 412fn capability_supported(json: &serde_json::Value, name: &str) -> bool {
413. 413 json.pointer(&format!("/capabilities/{name}/supported"))
414. 414 .and_then(|v| v.as_bool())
415. 415 .unwrap_or(false)
416. 416}
417. 417
418. 418/// Fetches one Anthropic model's capability flags from `GET /v1/models/{id}`
419. 419/// (invariant 9: discovered, never hardcoded — a static per-model-id table
420. 420/// would drift the moment a new model shipped or an old key lost access to
421. 421/// one). Meaningful for the `anthropic` provider only; other providers do
422. 422/// not publish this shape yet.
423. 423pub async fn anthropic_model_capabilities(
424. 424 auth: &ProviderAuth,
425. 425 model: &str,
426. 426) -> Result<AnthropicCapabilities, LlmError> {
427. 427 let base = auth
428. 428 .base_url
429. 429 .as_deref()
430. 430 .filter(|b| !b.trim().is_empty())
431. 431 .or_else(|| default_base_url("anthropic"))
432. 432 .ok_or_else(|| LlmError::InvalidRequest("no base url for 'anthropic'".into()))?
433. 433 .trim_end_matches('/')
434. 434 .to_string();
435. 435 let url = if base.ends_with("/v1") {
436. 436 format!("{base}/models/{model}")
437. 437 } else {
438. 438 format!("{base}/v1/models/{model}")
439. 439 };
440. 440 let response = http()?
441. 441 .get(&url)
442. 442 .header("x-api-key", &auth.api_key)
443. 443 .header("anthropic-version", crate::anthropic::ANTHROPIC_VERSION)
444. 444 .send()
445. 445 .await
446. 446 .map_err(|e| LlmError::Network(e.to_string()))?;
447. 447 let json = read_json(response).await?;
448. 448 Ok(AnthropicCapabilities {
449. 449 effort: capability_supported(&json, "effort"),
450. 450 fast_mode: capability_supported(&json, "fast_mode"),
451. 451 })
452. 452}
453. 453
454. 454/// In-process memory of what this run has learned about each Anthropic
455. 455/// model's `effort`/fast-mode support — never persisted, never shared
456. 456/// across processes, and only ever narrows what the adapter attempts next
457. 457/// (docs/design/68-context-engine.md §11). Two different policies live
458. 458/// behind the same shape because the two features start from opposite
459. 459/// priors: `effort` is GA on every current-generation model, so a model not
460. 460/// yet seen defaults to "try it" and a live 400 naming the parameter is
461. 461/// what teaches the cache `false`; fast mode is a narrow, opt-in research
462. 462/// preview, so a model not yet seen defaults to "do not try it" until an
463. 463/// explicit capability fetch has said otherwise.
464. 464#[derive(Default)]
465. 465struct AnthropicCapabilityCache {
466. 466 effort: std::collections::HashMap<String, bool>,
467. 467 fast_mode: std::collections::HashMap<String, bool>,
468. 468}
469. 469
470. 470fn anthropic_capability_cache() -> &'static std::sync::Mutex<AnthropicCapabilityCache> {
471. 471 static CACHE: std::sync::OnceLock<std::sync::Mutex<AnthropicCapabilityCache>> =
472. 472 std::sync::OnceLock::new();
473. 473 CACHE.get_or_init(|| std::sync::Mutex::new(AnthropicCapabilityCache::default()))
474. 474}
475. 475
476. 476/// Whether the adapter should attempt `output_config.effort` for `model`.
477. 477/// Defaults to `true` (unknown models are worth trying) until
478. 478/// [`mark_anthropic_effort_unsupported`] narrows it.
479. 479pub fn anthropic_effort_allowed(model: &str) -> bool {
480. 480 anthropic_capability_cache()
481. 481 .lock()
482. 482 .unwrap_or_else(std::sync::PoisonError::into_inner)
483. 483 .effort
484. 484 .get(model)
485. 485 .copied()
486. 486 .unwrap_or(true)
487. 487}
488. 488
489. 489/// Records that `model` rejected `output_config.effort` (a live 400 naming
490. 490/// the parameter): every later request in this process skips sending it.
491. 491pub fn mark_anthropic_effort_unsupported(model: &str) {
492. 492 anthropic_capability_cache()
493. 493 .lock()
494. 494 .unwrap_or_else(std::sync::PoisonError::into_inner)
495. 495 .effort
496. 496 .insert(model.to_string(), false);
497. 497}
498. 498
499. 499/// Whether the adapter should attempt `speed: "fast"` for `model`. Defaults
500. 500/// to `false` (an unknown model is never assumed to support a research
501. 501/// preview) until [`record_anthropic_capabilities`] has looked it up.
502. 502pub fn anthropic_fast_mode_allowed(model: &str) -> bool {
503. 503 anthropic_capability_cache()
504. 504 .lock()
505. 505 .unwrap_or_else(std::sync::PoisonError::into_inner)
506. 506 .fast_mode
507. 507 .get(model)
508. 508 .copied()
509. 509 .unwrap_or(false)
510. 510}
511. 511
512. 512/// Whether `model`'s fast-mode support has already been learned one way or
513. 513/// the other, so a caller can skip a redundant discovery fetch.
514. 514pub fn anthropic_fast_mode_known(model: &str) -> bool {
515. 515 anthropic_capability_cache()
516. 516 .lock()
517. 517 .unwrap_or_else(std::sync::PoisonError::into_inner)
518. 518 .fast_mode
519. 519 .contains_key(model)
520. 520}
521. 521
522. 522/// Records a freshly discovered capability set
523. 523/// ([`anthropic_model_capabilities`]) so later requests for `model` skip
524. 524/// both the network round trip and the conservative first-attempt default.
525. 525pub fn record_anthropic_capabilities(model: &str, caps: AnthropicCapabilities) {
526. 526 let mut cache = anthropic_capability_cache()
527. 527 .lock()
528. 528 .unwrap_or_else(std::sync::PoisonError::into_inner);
529. 529 cache.effort.insert(model.to_string(), caps.effort);
530. 530 cache.fast_mode.insert(model.to_string(), caps.fast_mode);
531. 531}
532. 532
533. 533/// The provider's documented API host, for callers that never set an
534. 534/// override. These are endpoints, not a model catalogue — the model list
535. 535/// itself always comes off the wire.
536. 536fn default_base_url(provider: &str) -> Option<&'static str> {
537. 537 match provider {
538. 538 "anthropic" => Some(crate::anthropic::DEFAULT_BASE_URL),
539. 539 "openai" => Some(crate::openai::OPENAI_DEFAULT_BASE_URL),
540. 540 "openai-responses" => Some(crate::openai_responses::OPENAI_RESPONSES_DEFAULT_BASE_URL),
541. 541 "google" => Some(crate::google::GOOGLE_DEFAULT_BASE_URL),
542. 542 "openrouter" => Some("https://openrouter.ai/api/v1"),
543. 543 "openrouter-responses" => Some("https://openrouter.ai/api/v1"),
544. 544 "bedrock" => Some("https://bedrock-mantle.us-east-1.api.aws/v1"),
545. 545 "opencode-zen" => Some("https://opencode.ai/zen/v1"),
546. 546 "ollama" => Some("http://localhost:11434/v1"),
547. 547 _ => None,
548. 548 }
549. 549}
550. 550
551. 551/// Google returns fully qualified `models/gemini-x`; the chat path wants
552. 552/// the bare id.
553. 553fn collect_google_names(json: &serde_json::Value) -> Vec<String> {
554. 554 json.get("models")
555. 555 .and_then(|m| m.as_array())
556. 556 .map(|arr| {
557. 557 arr.iter()
558. 558 .filter_map(|m| m.get("name").and_then(|n| n.as_str()))
559. 559 .map(|n| n.trim_start_matches("models/").to_string())
560. 560 .collect()
561. 561 })
562. 562 .unwrap_or_default()
563. 563}
564. 564
565. 565fn collect_data_ids(json: serde_json::Value) -> Vec<String> {
566. 566 json.get("data")
567. 567 .and_then(|d| d.as_array())
568. 568 .map(|arr| {
569. 569 arr.iter()
570. 570 .filter_map(|m| m.get("id").and_then(|i| i.as_str()))
571. 571 .map(str::to_string)
572. 572 .collect()
573. 573 })
574. 574 .unwrap_or_default()
575. 575}
576. 576
577. 577#[cfg(test)]
578. 578mod tests {
579. 579 use super::*;
580. 580
581. 581 #[test]
582. 582 fn openai_shape_yields_ids() {
583. 583 let json = serde_json::json!({ "data": [{ "id": "gpt-4.1" }, { "id": "o3" }] });
584. 584 assert_eq!(collect_data_ids(json), vec!["gpt-4.1", "o3"]);
585. 585 }
586. 586
587. 587 #[test]
588. 588 fn missing_data_is_still_empty_for_the_parser() {
589. 589 assert!(collect_data_ids(serde_json::json!({})).is_empty());
590. 590 }
591. 591
592. 592 #[test]
593. 593 fn google_names_are_stripped_of_the_models_prefix() {
594. 594 let json = serde_json::json!({ "models": [{ "name": "models/gemini-2.5-pro" }] });
595. 595 assert_eq!(collect_google_names(&json), vec!["gemini-2.5-pro"]);
596. 596 }
597. 597
598. 598 #[test]
599. 599 fn context_metadata_prefers_provider_specific_limit() {
600. 600 let json = serde_json::json!({
601. 601 "data": {"context_length": 131072, "top_provider": {"context_length": 65536, "max_completion_tokens": 8192}}
602. 602 });
603. 603 assert_eq!(
604. 604 context_from_json("openrouter", &json),
605. 605 Some(ModelContext {
606. 606 input_tokens: 65536,
607. 607 output_tokens: Some(8192),
608. 608 quantisation: None,
609. 609 })
610. 610 );
611. 611 }
612. 612
613. 613 #[test]
614. 614 fn every_supported_provider_has_a_default_endpoint() {
615. 615 for p in [
616. 616 "anthropic",
617. 617 "openai",
618. 618 "openai-responses",
619. 619 "google",
620. 620 "openrouter",
621. 621 "openrouter-responses",
622. 622 "opencode-zen",
623. 623 "ollama",
624. 624 ] {
625. 625 assert!(default_base_url(p).is_some(), "{p} has no default base url");
626. 626 }
627. 627 }
628. 628
629. 629 #[test]
630. 630 fn ollama_context_extracted_from_model_info() {
631. 631 let json = serde_json::json!({
632. 632 "model_info": {
633. 633 "qwen35.context_length": 32768
634. 634 }
635. 635 });
636. 636 assert_eq!(
637. 637 context_from_json("ollama", &json),
638. 638 Some(ModelContext {
639. 639 input_tokens: 32768,
640. 640 output_tokens: Some(4096),
641. 641 quantisation: None,
642. 642 })
643. 643 );
644. 644 }
645. 645
646. 646 #[test]
647. 647 fn ollama_context_prefers_modelfile_num_ctx() {
648. 648 let json = serde_json::json!({
649. 649 "parameters": "temperature 0.7\nnum_ctx 16384\ntop_p 0.9",
650. 650 "model_info": {
651. 651 "general.context_length": 131072
652. 652 }
653. 653 });
654. 654 assert_eq!(
655. 655 context_from_json("ollama", &json),
656. 656 Some(ModelContext {
657. 657 input_tokens: 16384,
658. 658 output_tokens: Some(4096),
659. 659 quantisation: None,
660. 660 })
661. 661 );
662. 662 }
663. 663
664. 664 #[test]
665. 665 fn ollama_context_parses_quantisation_from_details() {
666. 666 let json = serde_json::json!({
667. 667 "details": {"quantization_level": "Q4_K_M"},
668. 668 "model_info": {"general.context_length": 8192}
669. 669 });
670. 670 assert_eq!(
671. 671 context_from_json("ollama", &json),
672. 672 Some(ModelContext {
673. 673 input_tokens: 8192,
674. 674 output_tokens: Some(4096),
675. 675 quantisation: Some("Q4_K_M".to_string()),
676. 676 })
677. 677 );
678. 678 }
679. 679
680. 680 #[test]
681. 681 fn unauthorised_maps_to_auth_error() {
682. 682 assert!(matches!(
683. 683 status_error(401, "bad key".into()),
684. 684 LlmError::Auth(_)
685. 685 ));
686. 686 }
687. 687
688. 688 #[test]
689. 689 fn capability_supported_reads_the_nested_flag() {
690. 690 let json = serde_json::json!({
691. 691 "capabilities": {
692. 692 "effort": {"supported": true},
693. 693 "thinking": {"supported": false},
694. 694 }
695. 695 });
696. 696 assert!(capability_supported(&json, "effort"));
697. 697 assert!(!capability_supported(&json, "thinking"));
698. 698 // Absent / unrecognized keys read as false, never guessed true.
699. 699 assert!(!capability_supported(&json, "fast_mode"));
700. 700 assert!(!capability_supported(&serde_json::json!({}), "effort"));
701. 701 }
702. 702
703. 703 // Each test below uses its own unique model id: the capability cache is
704. 704 // a process-wide static, and tests in this module run concurrently.
705. 705 #[test]
706. 706 fn effort_defaults_to_allowed_for_an_unseen_model() {
707. 707 assert!(anthropic_effort_allowed("test-model-effort-unseen-1"));
708. 708 }
709. 709
710. 710 #[test]
711. 711 fn marking_effort_unsupported_narrows_it_for_that_model_only() {
712. 712 anthropic_effort_allowed("test-model-effort-sibling-2"); // establish baseline
713. 713 mark_anthropic_effort_unsupported("test-model-effort-marked-2");
714. 714 assert!(!anthropic_effort_allowed("test-model-effort-marked-2"));
715. 715 assert!(anthropic_effort_allowed("test-model-effort-sibling-2"));
716. 716 }
717. 717
718. 718 #[test]
719. 719 fn fast_mode_defaults_to_not_allowed_and_not_known_for_an_unseen_model() {
720. 720 assert!(!anthropic_fast_mode_allowed("test-model-fast-unseen-3"));
721. 721 assert!(!anthropic_fast_mode_known("test-model-fast-unseen-3"));
722. 722 }
723. 723
724. 724 #[test]
725. 725 fn recording_discovered_capabilities_makes_fast_mode_known_and_gates_on_the_flag() {
726. 726 record_anthropic_capabilities(
727. 727 "test-model-fast-supported-4",
728. 728 AnthropicCapabilities {
729. 729 effort: true,
730. 730 fast_mode: true,
731. 731 },
732. 732 );
733. 733 assert!(anthropic_fast_mode_known("test-model-fast-supported-4"));
734. 734 assert!(anthropic_fast_mode_allowed("test-model-fast-supported-4"));
735. 735
736. 736 record_anthropic_capabilities(
737. 737 "test-model-fast-unsupported-5",
738. 738 AnthropicCapabilities {
739. 739 effort: true,
740. 740 fast_mode: false,
741. 741 },
742. 742 );
743. 743 assert!(anthropic_fast_mode_known("test-model-fast-unsupported-5"));
744. 744 assert!(!anthropic_fast_mode_allowed(
745. 745 "test-model-fast-unsupported-5"
746. 746 ));
747. 747 }
748. 748}
749. 749