Files
KimiCodeSwitch/src-tauri/src/kimi_code_io.rs
T

531 lines
19 KiB
Rust

//! Kimi Code configuration I/O.
//!
//! Kimi Code stores its configuration in `~/.kimi-code/config.toml`.
//! This module reads/writes that file and converts between Kimi's TOML
//! format and Pi Switch's internal `Config`/`Provider`/`Model` types.
use std::path::PathBuf;
use anyhow::Context;
use indexmap::IndexMap;
use serde_json::Value;
use toml::value::{Table, Value as TomlValue};
use crate::config_io::backup_file;
use crate::models::{Agent, Config, Model, Provider, ProviderType};
pub type KimiResult<T> = anyhow::Result<T>;
pub fn kimi_code_config_dir() -> PathBuf {
std::env::var_os("KIMI_CODE_HOME")
.map(PathBuf::from)
.or_else(|| dirs::home_dir().map(|h| h.join(".kimi-code")))
.expect("failed to resolve Kimi Code config directory")
}
pub fn kimi_code_config_path() -> PathBuf {
kimi_code_config_dir().join("config.toml")
}
/// Read the existing Kimi Code `config.toml`, or return an empty table if it
/// does not exist yet.
pub fn load_kimi_code_config() -> KimiResult<TomlValue> {
let path = kimi_code_config_path();
if !path.exists() {
return Ok(TomlValue::Table(Table::new()));
}
let content = std::fs::read_to_string(&path)
.with_context(|| format!("failed to read {}", path.display()))?;
let value: TomlValue = content
.parse()
.with_context(|| format!("failed to parse {}", path.display()))?;
Ok(value)
}
/// Write the given TOML value to Kimi Code's config file, creating the
/// directory if needed and backing up the previous file.
pub fn save_kimi_code_config(value: &TomlValue) -> KimiResult<()> {
let path = kimi_code_config_path();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("failed to create directory {}", parent.display()))?;
}
if path.exists() {
backup_file(&path, 7).with_context(|| {
format!("failed to back up Kimi Code config {}", path.display())
})?;
}
let content = toml::to_string_pretty(value)
.with_context(|| "failed to serialize Kimi Code config.toml")?;
std::fs::write(&path, content)
.with_context(|| format!("failed to write {}", path.display()))?;
Ok(())
}
/// Import a Kimi Code TOML config into Pi Switch's internal `Config`.
pub fn kimi_code_to_config(value: &TomlValue) -> Config {
let root = value.as_table().cloned().unwrap_or_default();
let default_model = root
.get("default_model")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let mut providers = IndexMap::new();
let mut models = IndexMap::new();
if let Some(providers_table) = root.get("providers").and_then(|v| v.as_table()) {
for (name, pv) in providers_table {
let table = pv.as_table().cloned().unwrap_or_default();
let provider_type = table
.get("type")
.and_then(|v| v.as_str())
.map(provider_type_for_kimi_type)
.unwrap_or(ProviderType::Kimi);
let base_url = table.get("base_url").and_then(|v| v.as_str()).map(|s| s.to_string());
let api_key = table.get("api_key").and_then(|v| v.as_str()).map(|s| s.to_string());
let managed = table.contains_key("oauth")
|| table.get("managed").and_then(|v| v.as_bool()).unwrap_or(false);
let enabled = table.get("enabled").and_then(|v| v.as_bool()).unwrap_or(true);
let env: IndexMap<String, String> = table
.get("env")
.and_then(|v| v.as_table())
.map(|t| {
t.iter()
.filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
.collect()
})
.unwrap_or_default();
let raw_other = {
let mut rest = table.clone();
rest.remove("type");
rest.remove("base_url");
rest.remove("api_key");
rest.remove("managed");
rest.remove("enabled");
rest.remove("oauth");
rest.remove("env");
toml_value_to_json(&TomlValue::Table(rest))
};
providers.insert(
name.clone(),
Provider {
name: name.clone(),
provider_type,
base_url: base_url.filter(|s| !s.is_empty()),
api_key: api_key.filter(|s| !s.is_empty()),
env,
note: None,
official_url: None,
managed,
enabled,
active: true,
raw_other,
},
);
}
}
if let Some(models_table) = root.get("models").and_then(|v| v.as_table()) {
for (alias, mv) in models_table {
let table = mv.as_table().cloned().unwrap_or_default();
let provider = table
.get("provider")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_default();
let model_id = table
.get("model")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_default();
let max_context_size = table
.get("max_context_size")
.and_then(|v| v.as_integer())
.map(|n| n as u64)
.unwrap_or(128_000);
let display_name = table
.get("display_name")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let capabilities: Vec<String> = table
.get("capabilities")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(|s| s.to_string()))
.collect()
})
.unwrap_or_default();
let supports_1m = capabilities.iter().any(|c| c == "thinking")
|| max_context_size >= 1_000_000;
let raw_other = {
let mut rest = table.clone();
rest.remove("provider");
rest.remove("model");
rest.remove("max_context_size");
rest.remove("display_name");
rest.remove("capabilities");
toml_value_to_json(&TomlValue::Table(rest))
};
models.insert(
alias.clone(),
Model {
alias: alias.clone(),
provider,
model: model_id,
max_context_size,
display_name,
role: None,
supports_1m,
capabilities,
raw_other,
},
);
}
}
let raw_other = {
let mut rest = root;
rest.remove("default_model");
rest.remove("providers");
rest.remove("models");
toml_value_to_json(&TomlValue::Table(rest))
};
Config {
default_model,
providers,
models,
raw_other,
}
}
/// Export Pi Switch's internal `Config` to a Kimi Code TOML config value.
///
/// If `existing` is provided, unknown top-level sections (e.g. `services`) are
/// preserved; otherwise a fresh TOML table is used.
pub fn config_to_kimi_code(config: &Config, existing: Option<&TomlValue>) -> TomlValue {
let mut root = existing
.and_then(|v| v.as_table().cloned())
.unwrap_or_default();
root.insert(
"default_model".to_string(),
config
.default_model
.clone()
.map(TomlValue::String)
.unwrap_or(TomlValue::String("".to_string())),
);
let mut providers_table = Table::new();
for (name, provider) in &config.providers {
// Only write the active provider to Kimi Code config so the agent
// follows Pi Switch's selection.
if !provider.active {
continue;
}
let mut pt = Table::new();
pt.insert("type".to_string(), TomlValue::String(provider.provider_type.as_str().to_string()));
if let Some(base_url) = provider.base_url.clone().filter(|s| !s.is_empty()) {
pt.insert("base_url".to_string(), TomlValue::String(base_url));
}
if let Some(api_key) = provider.api_key.clone().filter(|s| !s.is_empty()) {
pt.insert("api_key".to_string(), TomlValue::String(api_key));
}
pt.insert("enabled".to_string(), TomlValue::Boolean(provider.enabled));
if !provider.env.is_empty() {
let mut env_table = Table::new();
for (k, v) in &provider.env {
env_table.insert(k.clone(), TomlValue::String(v.clone()));
}
pt.insert("env".to_string(), TomlValue::Table(env_table));
}
if provider.managed {
// Preserve existing oauth config if present, otherwise create a default entry.
let oauth = provider
.raw_other
.get("oauth")
.and_then(json_to_toml)
.unwrap_or_else(|| {
let mut t = Table::new();
t.insert("storage".to_string(), TomlValue::String("file".to_string()));
t.insert("key".to_string(), TomlValue::String(format!("oauth/{}", name.replace(':', "-"))));
TomlValue::Table(t)
});
pt.insert("oauth".to_string(), oauth);
}
// Merge remaining raw fields (oauth and env are handled explicitly above).
if let TomlValue::Table(mut extra) = json_to_toml(&provider.raw_other).unwrap_or(TomlValue::Table(Table::new())) {
extra.remove("oauth");
extra.remove("env");
for (k, v) in extra {
pt.insert(k, v);
}
}
providers_table.insert(name.clone(), TomlValue::Table(pt));
}
root.insert("providers".to_string(), TomlValue::Table(providers_table));
// Collect provider names that survived the filter above.
let active_provider_names: std::collections::HashSet<&str> = config
.providers
.values()
.filter(|p| p.active)
.map(|p| p.name.as_str())
.collect();
let mut models_table = Table::new();
for (alias, model) in &config.models {
if !active_provider_names.contains(model.provider.as_str()) {
continue;
}
let mut mt = Table::new();
mt.insert("provider".to_string(), TomlValue::String(model.provider.clone()));
mt.insert("model".to_string(), TomlValue::String(model.model.clone()));
mt.insert(
"max_context_size".to_string(),
TomlValue::Integer(model.max_context_size as i64),
);
if let Some(display_name) = model.display_name.clone().filter(|s| !s.is_empty()) {
mt.insert("display_name".to_string(), TomlValue::String(display_name));
}
let mut capabilities = model.capabilities.clone();
if model.supports_1m && !capabilities.iter().any(|c| c == "thinking") {
capabilities.push("thinking".to_string());
}
if !capabilities.is_empty() {
mt.insert(
"capabilities".to_string(),
TomlValue::Array(capabilities.into_iter().map(TomlValue::String).collect()),
);
}
if let TomlValue::Table(mut extra) = json_to_toml(&model.raw_other).unwrap_or(TomlValue::Table(Table::new())) {
extra.remove("provider");
extra.remove("model");
extra.remove("max_context_size");
extra.remove("display_name");
extra.remove("capabilities");
for (k, v) in extra {
mt.insert(k, v);
}
}
models_table.insert(alias.clone(), TomlValue::Table(mt));
}
root.insert("models".to_string(), TomlValue::Table(models_table));
TomlValue::Table(root)
}
fn provider_type_for_kimi_type(typ: &str) -> ProviderType {
match typ {
"anthropic" => ProviderType::Anthropic,
"openai" => ProviderType::Openai,
"openai_responses" => ProviderType::OpenaiResponses,
"google-genai" => ProviderType::GoogleGenai,
"vertexai" => ProviderType::Vertexai,
_ => ProviderType::Kimi,
}
}
fn toml_value_to_json(value: &TomlValue) -> Value {
match value {
TomlValue::String(s) => Value::String(s.clone()),
TomlValue::Integer(n) => Value::Number((*n).into()),
TomlValue::Float(n) => Value::Number(serde_json::Number::from_f64(*n).unwrap_or(Value::Null.as_u64().unwrap_or(0).into())),
TomlValue::Boolean(b) => Value::Bool(*b),
TomlValue::Array(arr) => Value::Array(arr.iter().map(toml_value_to_json).collect()),
TomlValue::Table(t) => {
let mut map = serde_json::Map::new();
for (k, v) in t {
map.insert(k.clone(), toml_value_to_json(v));
}
Value::Object(map)
}
TomlValue::Datetime(dt) => Value::String(dt.to_string()),
}
}
fn json_to_toml(value: &Value) -> Option<TomlValue> {
match value {
Value::Null => None,
Value::Bool(b) => Some(TomlValue::Boolean(*b)),
Value::Number(n) => {
if let Some(i) = n.as_i64() {
Some(TomlValue::Integer(i))
} else if let Some(f) = n.as_f64() {
Some(TomlValue::Float(f))
} else {
None
}
}
Value::String(s) => Some(TomlValue::String(s.clone())),
Value::Array(arr) => Some(TomlValue::Array(
arr.iter().filter_map(json_to_toml).collect(),
)),
Value::Object(obj) => {
let mut t = Table::new();
for (k, v) in obj {
if let Some(tv) = json_to_toml(v) {
t.insert(k.clone(), tv);
}
}
Some(TomlValue::Table(t))
}
}
}
/// Convenience entry point used by commands: load and convert.
pub fn load_kimi_code_config_as_config() -> KimiResult<Config> {
let value = load_kimi_code_config()?;
Ok(kimi_code_to_config(&value))
}
/// Convenience entry point used by commands: convert and save.
pub fn save_config_as_kimi_code(config: &Config) -> KimiResult<()> {
let existing = load_kimi_code_config().ok();
let value = config_to_kimi_code(config, existing.as_ref());
save_kimi_code_config(&value)
}
/// Open the Kimi Code config directory in the system file manager.
pub fn open_kimi_code_config_dir() -> KimiResult<()> {
let path = kimi_code_config_dir();
std::fs::create_dir_all(&path)
.with_context(|| format!("failed to create directory {}", path.display()))?;
Ok(())
}
impl Agent {
pub fn config_dir(&self) -> PathBuf {
match self {
Agent::KimiCode => kimi_code_config_dir(),
Agent::Pi => crate::pi_io::pi_agent_dir(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn kimi_code_import_extracts_providers_and_models() {
let toml_str = r#"
default_model = "kimi-code/kimi-for-coding"
default_thinking = true
[providers."managed:kimi-code"]
type = "kimi"
api_key = ""
base_url = "https://api.kimi.com/coding/v1"
[providers."managed:kimi-code".oauth]
storage = "file"
key = "oauth/kimi-code"
[models."kimi-code/kimi-for-coding"]
provider = "managed:kimi-code"
model = "kimi-for-coding"
max_context_size = 262144
capabilities = ["thinking", "always_thinking", "image_in", "video_in", "tool_use"]
display_name = "K2.7 Code"
[services.moonshot_search]
base_url = "https://api.kimi.com/coding/v1/search"
api_key = ""
"#;
let value: TomlValue = toml_str.parse().unwrap();
let config = kimi_code_to_config(&value);
assert_eq!(config.default_model.as_deref(), Some("kimi-code/kimi-for-coding"));
assert_eq!(config.providers.len(), 1);
assert_eq!(config.models.len(), 1);
let provider = config.providers.get("managed:kimi-code").unwrap();
assert_eq!(provider.provider_type, ProviderType::Kimi);
assert!(provider.managed);
assert_eq!(provider.base_url.as_deref(), Some("https://api.kimi.com/coding/v1"));
let model = config.models.get("kimi-code/kimi-for-coding").unwrap();
assert_eq!(model.provider, "managed:kimi-code");
assert_eq!(model.model, "kimi-for-coding");
assert_eq!(model.max_context_size, 262144);
assert!(model.capabilities.contains(&"thinking".to_string()));
assert!(model.supports_1m);
// Unknown top-level sections are preserved.
let services = config.raw_other.get("services");
assert!(services.is_some());
}
#[test]
fn kimi_code_export_roundtrip_preserves_model() {
let mut providers = IndexMap::new();
providers.insert(
"glmzhongzhuan".to_string(),
Provider {
name: "glmzhongzhuan".to_string(),
provider_type: ProviderType::Anthropic,
base_url: Some("https://fast.cdks.work".to_string()),
api_key: Some("sk-test".to_string()),
env: IndexMap::new(),
note: None,
official_url: None,
managed: false,
enabled: true,
active: true,
raw_other: Value::Null,
},
);
let mut models = IndexMap::new();
models.insert(
"glm-5.2".to_string(),
Model {
alias: "glm-5.2".to_string(),
provider: "glmzhongzhuan".to_string(),
model: "glm-5.2".to_string(),
max_context_size: 900_000,
display_name: None,
role: None,
supports_1m: true,
capabilities: vec!["thinking".to_string()],
raw_other: Value::Null,
},
);
let config = Config {
default_model: Some("glm-5.2".to_string()),
providers,
models,
raw_other: Value::Null,
};
let exported = config_to_kimi_code(&config, None);
let root = exported.as_table().unwrap();
assert_eq!(
root.get("default_model").and_then(|v| v.as_str()),
Some("glm-5.2")
);
let providers_table = root.get("providers").unwrap().as_table().unwrap();
let provider = providers_table.get("glmzhongzhuan").unwrap().as_table().unwrap();
assert_eq!(provider.get("type").and_then(|v| v.as_str()), Some("anthropic"));
assert_eq!(provider.get("api_key").and_then(|v| v.as_str()), Some("sk-test"));
let models_table = root.get("models").unwrap().as_table().unwrap();
let model = models_table.get("glm-5.2").unwrap().as_table().unwrap();
assert_eq!(model.get("model").and_then(|v| v.as_str()), Some("glm-5.2"));
let caps = model.get("capabilities").unwrap().as_array().unwrap();
assert!(caps.iter().any(|v| v.as_str() == Some("thinking")));
}
}