code cleaning, toml edit helpers

This commit is contained in:
pap
2025-08-04 00:37:18 +01:00
parent 75eec73fcc
commit b2ed15430f
6 changed files with 105 additions and 122 deletions

8
codex-rs/Cargo.lock generated
View File

@@ -709,6 +709,7 @@ dependencies = [
"tokio-test",
"tokio-util",
"toml 0.9.2",
"toml_edit",
"tracing",
"tree-sitter",
"tree-sitter-bash",
@@ -4815,6 +4816,7 @@ dependencies = [
"serde",
"serde_spanned 0.6.9",
"toml_datetime 0.6.11",
"toml_write",
"winnow",
]
@@ -4827,6 +4829,12 @@ dependencies = [
"winnow",
]
[[package]]
name = "toml_write"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801"
[[package]]
name = "toml_writer"
version = "1.0.2"

View File

@@ -48,6 +48,7 @@ tokio = { version = "1", features = [
] }
tokio-util = "0.7.14"
toml = "0.9.2"
toml_edit = "0.22"
tracing = { version = "0.1.41", features = ["log"] }
tree-sitter = "0.25.8"
tree-sitter-bash = "0.25.0"

View File

@@ -37,6 +37,7 @@ mod openai_tools;
pub mod plan_tool;
mod project_doc;
pub mod protocol;
pub mod providers;
mod rollout;
pub(crate) mod safety;
pub mod seatbelt;
@@ -45,7 +46,6 @@ pub mod spawn;
pub mod turn_diff_tracker;
mod user_notification;
pub mod util;
pub mod providers;
pub use apply_patch::CODEX_APPLY_PATCH_ARG1;
pub use client_common::model_supports_reasoning_summaries;

View File

@@ -1,2 +1 @@
pub mod ollama;

View File

@@ -9,6 +9,11 @@ use std::collections::VecDeque;
use std::io;
use std::io::Write;
use std::path::Path;
use std::str::FromStr;
use toml_edit::DocumentMut as Document;
use toml_edit::Item;
use toml_edit::Table;
use toml_edit::Value as TomlValueEdit;
pub const DEFAULT_BASE_URL: &str = "http://localhost:11434/v1";
pub const DEFAULT_WIRE_API: WireApi = WireApi::Chat;
@@ -33,6 +38,12 @@ pub fn base_url_to_host_root(base_url: &str) -> String {
}
}
/// Variant that considers an explicit WireApi value; provided to centralize
/// host root computation in one place for future extension.
pub fn base_url_to_host_root_with_wire(base_url: &str, _wire_api: WireApi) -> String {
base_url_to_host_root(base_url)
}
/// Compute the probe URL to verify if an Ollama server is reachable.
/// If the configured base is OpenAI-compatible (/v1), probe "models", otherwise
/// fall back to the native "/api/tags" endpoint.
@@ -46,9 +57,9 @@ pub fn probe_url_for_base(base_url: &str) -> String {
/// Convenience helper to probe an Ollama server given a provider style base URL.
pub async fn probe_ollama_server(base_url: &str) -> io::Result<bool> {
let host_root = base_url_to_host_root(base_url);
let client = OllamaClient::from_host_root(host_root);
client.probe_server().await
let url = probe_url_for_base(base_url);
let resp = reqwest::Client::new().get(url).send().await;
Ok(matches!(resp, Ok(r) if r.status().is_success()))
}
/// Coordinator wrapper used by frontends when responding to `--ollama`.
///
@@ -423,58 +434,18 @@ pub fn read_config_models(config_path: &Path) -> io::Result<Vec<String>> {
Ok(read_ollama_models_list(config_path))
}
/// Overwrite the recorded models list under [model_providers.ollama].models.
/// Overwrite the recorded models list under [model_providers.ollama].models using toml_edit.
pub fn write_ollama_models_list(config_path: &Path, models: &[String]) -> io::Result<()> {
use toml::value::Table as TomlTable;
let mut root_value = if let Ok(contents) = std::fs::read_to_string(config_path) {
toml::from_str::<toml::Value>(&contents).unwrap_or(toml::Value::Table(TomlTable::new()))
} else {
toml::Value::Table(TomlTable::new())
};
if !matches!(root_value, toml::Value::Table(_)) {
root_value = toml::Value::Table(TomlTable::new());
let mut doc = read_document(config_path)?;
{
let tbl = upsert_provider_ollama(&mut doc);
let mut arr = toml_edit::Array::new();
for m in models {
arr.push(TomlValueEdit::from(m.clone()));
}
tbl["models"] = Item::Value(TomlValueEdit::Array(arr));
}
let root_tbl = match root_value.as_table_mut() {
Some(t) => t,
None => return Err(io::Error::other("invalid TOML root value")),
};
let mp_val = root_tbl
.entry("model_providers".to_string())
.or_insert_with(|| toml::Value::Table(TomlTable::new()));
if !mp_val.is_table() {
*mp_val = toml::Value::Table(TomlTable::new());
}
let mp_tbl = match mp_val.as_table_mut() {
Some(t) => t,
None => return Err(io::Error::other("invalid model_providers table")),
};
let ollama_val = mp_tbl
.entry("ollama".to_string())
.or_insert_with(|| toml::Value::Table(TomlTable::new()));
if !ollama_val.is_table() {
*ollama_val = toml::Value::Table(TomlTable::new());
}
let ollama_tbl = match ollama_val.as_table_mut() {
Some(t) => t,
None => return Err(io::Error::other("invalid ollama table")),
};
let arr = toml::Value::Array(
models
.iter()
.map(|m| toml::Value::String(m.clone()))
.collect(),
);
ollama_tbl.insert("models".to_string(), arr);
let updated =
toml::to_string_pretty(&root_value).map_err(|e| io::Error::other(e.to_string()))?;
if let Some(parent) = config_path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(config_path, updated)
write_document(config_path, &doc)
}
/// Write models list via a uniform name expected by higher layers.
@@ -485,74 +456,17 @@ pub fn write_config_models(config_path: &Path, models: &[String]) -> io::Result<
/// Ensure `[model_providers.ollama]` exists with sensible defaults on disk.
/// Returns true if it created/updated the entry.
pub fn ensure_ollama_provider_entry(codex_home: &Path) -> io::Result<bool> {
use toml::value::Table as TomlTable;
let config_path = codex_home.join("config.toml");
let mut root_value = if let Ok(contents) = std::fs::read_to_string(&config_path) {
toml::from_str::<toml::Value>(&contents).unwrap_or(toml::Value::Table(TomlTable::new()))
let mut doc = read_document(&config_path)?;
let before = doc.to_string();
let _tbl = upsert_provider_ollama(&mut doc);
let after = doc.to_string();
if before != after {
write_document(&config_path, &doc)?;
Ok(true)
} else {
toml::Value::Table(TomlTable::new())
};
if !matches!(root_value, toml::Value::Table(_)) {
root_value = toml::Value::Table(TomlTable::new());
Ok(false)
}
let root_tbl = match root_value.as_table_mut() {
Some(t) => t,
None => return Err(io::Error::other("invalid TOML root")),
};
let mp_val = root_tbl
.entry("model_providers".to_string())
.or_insert_with(|| toml::Value::Table(TomlTable::new()));
if !mp_val.is_table() {
*mp_val = toml::Value::Table(TomlTable::new());
}
let mp_tbl = match mp_val.as_table_mut() {
Some(t) => t,
None => return Err(io::Error::other("invalid model_providers table")),
};
let mut changed = false;
let ollama_val = mp_tbl.entry("ollama".to_string()).or_insert_with(|| {
changed = true;
toml::Value::Table(TomlTable::new())
});
if !ollama_val.is_table() {
*ollama_val = toml::Value::Table(TomlTable::new());
changed = true;
}
if let Some(tbl) = ollama_val.as_table_mut() {
if !tbl.contains_key("name") {
tbl.insert(
"name".to_string(),
toml::Value::String("Ollama".to_string()),
);
changed = true;
}
if !tbl.contains_key("base_url") {
tbl.insert(
"base_url".to_string(),
toml::Value::String(DEFAULT_BASE_URL.to_string()),
);
changed = true;
}
if !tbl.contains_key("wire_api") {
tbl.insert(
"wire_api".to_string(),
toml::Value::String("chat".to_string()),
);
changed = true;
}
}
if changed {
if let Some(parent) = config_path.parent() {
std::fs::create_dir_all(parent)?;
}
let updated = toml::to_string_pretty(&root_value).map_err(io::Error::other)?;
std::fs::write(config_path, updated)?;
}
Ok(changed)
}
/// Alias name mirroring the refactor plan wording.
@@ -635,3 +549,61 @@ pub async fn ensure_model_available(
let _ = write_ollama_models_list(config_path, &listed);
Ok(())
}
// ---------- toml_edit helpers ----------
fn read_document(path: &Path) -> io::Result<Document> {
match std::fs::read_to_string(path) {
Ok(s) => Document::from_str(&s).map_err(io::Error::other),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Document::new()),
Err(e) => Err(e),
}
}
fn write_document(path: &Path, doc: &Document) -> io::Result<()> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, doc.to_string())
}
pub fn upsert_provider_ollama(doc: &mut Document) -> &mut Table {
// Ensure the provider tables exist first, then take a single mutable borrow.
if !doc["model_providers"].is_table() {
doc["model_providers"] = Item::Table(Table::new());
}
{
// Narrow scope: mutate/create the nested `ollama` table without keeping a borrow alive.
let providers = match doc["model_providers"].as_table_mut() {
Some(table) => table,
None => return Box::leak(Box::new(Table::default())),
};
if !providers.contains_key("ollama") || !providers["ollama"].is_table() {
providers["ollama"] = Item::Table(Table::new());
}
}
// Now, safely borrow the `ollama` table mutably once and return it.
let tbl = match doc["model_providers"]["ollama"].as_table_mut() {
Some(table) => table,
None => return Box::leak(Box::new(Table::default())),
};
if !tbl.contains_key("name") {
tbl["name"] = Item::Value(TomlValueEdit::from("Ollama"));
}
if !tbl.contains_key("base_url") {
tbl["base_url"] = Item::Value(TomlValueEdit::from(DEFAULT_BASE_URL));
}
if !tbl.contains_key("wire_api") {
tbl["wire_api"] = Item::Value(TomlValueEdit::from("chat"));
}
tbl
}
pub fn set_ollama_models(doc: &mut Document, models: &[String]) {
let tbl = upsert_provider_ollama(doc);
let mut arr = toml_edit::Array::new();
for m in models {
arr.push(TomlValueEdit::from(m.clone()));
}
tbl["models"] = Item::Value(TomlValueEdit::Array(arr));
}

View File

@@ -58,8 +58,11 @@ pub use cli::Cli;
// Centralized Ollama helpers from core
use codex_core::providers::ollama::OllamaClient;
use codex_core::providers::ollama::TuiProgressReporter;
use codex_core::providers::ollama::{ensure_configured_and_running, ensure_model_available};
use codex_core::providers::ollama::{read_config_models, read_provider_state, write_config_models};
use codex_core::providers::ollama::ensure_configured_and_running;
use codex_core::providers::ollama::ensure_model_available;
use codex_core::providers::ollama::read_config_models;
use codex_core::providers::ollama::read_provider_state;
use codex_core::providers::ollama::write_config_models;
fn print_inline_message_no_models(
host_root: &str,