|
1 | 1 | use bytes::BytesMut; |
2 | 2 | use futures::StreamExt; |
3 | 3 | use futures::stream::BoxStream; |
| 4 | +use semver::Version; |
4 | 5 | use serde_json::Value as JsonValue; |
5 | 6 | use std::collections::VecDeque; |
6 | 7 | use std::io; |
@@ -53,7 +54,7 @@ impl OllamaClient { |
53 | 54 | } |
54 | 55 |
|
55 | 56 | /// Build a client from a provider definition and verify the server is reachable. |
56 | | - async fn try_from_provider(provider: &ModelProviderInfo) -> io::Result<Self> { |
| 57 | + pub(crate) async fn try_from_provider(provider: &ModelProviderInfo) -> io::Result<Self> { |
57 | 58 | #![expect(clippy::expect_used)] |
58 | 59 | let base_url = provider |
59 | 60 | .base_url |
@@ -125,6 +126,32 @@ impl OllamaClient { |
125 | 126 | Ok(names) |
126 | 127 | } |
127 | 128 |
|
| 129 | + /// Query the server for its version string, returning `None` when unavailable. |
| 130 | + pub async fn fetch_version(&self) -> io::Result<Option<Version>> { |
| 131 | + let version_url = format!("{}/api/version", self.host_root.trim_end_matches('/')); |
| 132 | + let resp = self |
| 133 | + .client |
| 134 | + .get(version_url) |
| 135 | + .send() |
| 136 | + .await |
| 137 | + .map_err(io::Error::other)?; |
| 138 | + if !resp.status().is_success() { |
| 139 | + return Ok(None); |
| 140 | + } |
| 141 | + let val = resp.json::<JsonValue>().await.map_err(io::Error::other)?; |
| 142 | + let Some(version_str) = val.get("version").and_then(|v| v.as_str()).map(str::trim) else { |
| 143 | + return Ok(None); |
| 144 | + }; |
| 145 | + let normalized = version_str.trim_start_matches('v'); |
| 146 | + match Version::parse(normalized) { |
| 147 | + Ok(version) => Ok(Some(version)), |
| 148 | + Err(err) => { |
| 149 | + tracing::warn!("Failed to parse Ollama version `{version_str}`: {err}"); |
| 150 | + Ok(None) |
| 151 | + } |
| 152 | + } |
| 153 | + } |
| 154 | + |
128 | 155 | /// Start a model pull and emit streaming events. The returned stream ends when |
129 | 156 | /// a Success event is observed or the server closes the connection. |
130 | 157 | pub async fn pull_model_stream( |
@@ -236,6 +263,7 @@ impl OllamaClient { |
236 | 263 | #[cfg(test)] |
237 | 264 | mod tests { |
238 | 265 | use super::*; |
| 266 | + use pretty_assertions::assert_eq; |
239 | 267 |
|
240 | 268 | // Happy-path tests using a mock HTTP server; skip if sandbox network is disabled. |
241 | 269 | #[tokio::test] |
@@ -269,6 +297,42 @@ mod tests { |
269 | 297 | assert!(models.contains(&"mistral".to_string())); |
270 | 298 | } |
271 | 299 |
|
| 300 | + #[tokio::test] |
| 301 | + async fn test_fetch_version() { |
| 302 | + if std::env::var(codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() { |
| 303 | + tracing::info!( |
| 304 | + "{} is set; skipping test_fetch_version", |
| 305 | + codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR |
| 306 | + ); |
| 307 | + return; |
| 308 | + } |
| 309 | + |
| 310 | + let server = wiremock::MockServer::start().await; |
| 311 | + wiremock::Mock::given(wiremock::matchers::method("GET")) |
| 312 | + .and(wiremock::matchers::path("/api/tags")) |
| 313 | + .respond_with(wiremock::ResponseTemplate::new(200).set_body_raw( |
| 314 | + serde_json::json!({ "models": [] }).to_string(), |
| 315 | + "application/json", |
| 316 | + )) |
| 317 | + .mount(&server) |
| 318 | + .await; |
| 319 | + wiremock::Mock::given(wiremock::matchers::method("GET")) |
| 320 | + .and(wiremock::matchers::path("/api/version")) |
| 321 | + .respond_with(wiremock::ResponseTemplate::new(200).set_body_raw( |
| 322 | + serde_json::json!({ "version": "0.14.1" }).to_string(), |
| 323 | + "application/json", |
| 324 | + )) |
| 325 | + .mount(&server) |
| 326 | + .await; |
| 327 | + |
| 328 | + let client = OllamaClient::try_from_provider_with_base_url(server.uri().as_str()) |
| 329 | + .await |
| 330 | + .expect("client"); |
| 331 | + |
| 332 | + let version = client.fetch_version().await.expect("version fetch"); |
| 333 | + assert_eq!(version, Some(Version::new(0, 14, 1))); |
| 334 | + } |
| 335 | + |
272 | 336 | #[tokio::test] |
273 | 337 | async fn test_probe_server_happy_path_openai_compat_and_native() { |
274 | 338 | if std::env::var(codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() { |
|
0 commit comments