use serde::{Deserialize, Serialize};
use crate::{KIMETSU_CONFIG_VERSION, KimetsuResult};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProjectConfig {
pub kimetsu: KimetsuSection,
pub model: ModelSection,
pub broker: BrokerSection,
pub shell: ShellSection,
pub ingestion: IngestionSection,
pub run: RunSection,
#[serde(default)]
pub embedder: EmbedderSection,
#[serde(default)]
pub learning: LearningSection,
}
impl ProjectConfig {
pub fn default_for_project(project_id: impl Into<String>) -> Self {
Self {
kimetsu: KimetsuSection {
project_id: project_id.into(),
schema_version: KIMETSU_CONFIG_VERSION,
use_user_brain: default_true(),
mcp_write_tools: default_true(),
},
model: ModelSection::default(),
broker: BrokerSection::default(),
shell: ShellSection::default(),
ingestion: IngestionSection::default(),
run: RunSection::default(),
embedder: EmbedderSection::default(),
learning: LearningSection::default(),
}
}
pub fn from_toml(value: &str) -> KimetsuResult<Self> {
Ok(toml::from_str(value)?)
}
pub fn to_toml(&self) -> KimetsuResult<String> {
Ok(toml::to_string_pretty(self)?)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KimetsuSection {
pub project_id: String,
pub schema_version: i64,
#[serde(default = "default_true")]
pub use_user_brain: bool,
#[serde(default = "default_true")]
pub mcp_write_tools: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbedderSection {
#[serde(default = "default_embedder_id")]
pub model: String,
#[serde(default = "default_true")]
pub enabled: bool,
#[serde(default = "default_true")]
pub daemon: bool,
#[serde(default = "default_true")]
pub warm_on_start: bool,
#[serde(default = "default_reranker_id")]
pub reranker: String,
}
fn default_embedder_id() -> String {
"jina-v2-base-code".to_string()
}
fn default_reranker_id() -> String {
"ms-marco-tinybert-l-2-v2".to_string()
}
fn default_true() -> bool {
true
}
impl Default for EmbedderSection {
fn default() -> Self {
Self {
model: default_embedder_id(),
enabled: default_true(),
daemon: default_true(),
warm_on_start: default_true(),
reranker: default_reranker_id(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LearningSection {
#[serde(default = "default_auto_harvest")]
pub auto_harvest: bool,
#[serde(default = "default_true")]
pub store_queries: bool,
#[serde(default)]
pub distiller: DistillerSection,
}
fn default_auto_harvest() -> bool {
true
}
impl Default for LearningSection {
fn default() -> Self {
Self {
auto_harvest: default_auto_harvest(),
store_queries: default_true(),
distiller: DistillerSection::default(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistillerSection {
#[serde(default)]
pub enabled: bool,
#[serde(default = "default_distiller_provider")]
pub provider: String,
#[serde(default = "default_distiller_model")]
pub model: String,
#[serde(default = "default_distiller_api_key_env")]
pub api_key_env: String,
#[serde(default = "default_distiller_base_url_env")]
pub base_url_env: String,
#[serde(default)]
pub region: Option<String>,
#[serde(default = "default_distiller_region_env")]
pub region_env: String,
}
fn default_distiller_provider() -> String {
"anthropic".to_string()
}
fn default_distiller_model() -> String {
"claude-haiku-4-5".to_string()
}
fn default_distiller_api_key_env() -> String {
"ANTHROPIC_API_KEY".to_string()
}
fn default_distiller_base_url_env() -> String {
"ANTHROPIC_BASE_URL".to_string()
}
fn default_distiller_region_env() -> String {
"AWS_REGION".to_string()
}
impl Default for DistillerSection {
fn default() -> Self {
Self {
enabled: false,
provider: default_distiller_provider(),
model: default_distiller_model(),
api_key_env: default_distiller_api_key_env(),
base_url_env: default_distiller_base_url_env(),
region: None,
region_env: default_distiller_region_env(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelSection {
pub provider: String,
pub model: String,
pub api_key_env: String,
pub max_output_tokens: u32,
pub temperature: f32,
pub request_timeout_secs: u64,
#[serde(default)]
pub region: Option<String>,
#[serde(default = "default_region_env")]
pub region_env: String,
#[serde(default)]
pub price_per_mtok: Option<f64>,
}
fn default_region_env() -> String {
"AWS_REGION".to_string()
}
impl Default for ModelSection {
fn default() -> Self {
Self {
provider: "anthropic".to_string(),
model: "claude-opus-4-7".to_string(),
api_key_env: "ANTHROPIC_API_KEY".to_string(),
max_output_tokens: 8192,
temperature: 0.2,
request_timeout_secs: 120,
region: None,
region_env: default_region_env(),
price_per_mtok: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrokerSection {
pub default_budget_tokens: u32,
pub weights: BrokerWeights,
#[serde(default = "default_max_capsules")]
pub max_capsules: usize,
#[serde(default = "default_min_semantic_score")]
pub min_semantic_score: f32,
#[serde(default = "default_min_lexical_coverage")]
pub min_lexical_coverage: f32,
#[serde(default = "default_budget_floor_tokens")]
pub budget_floor_tokens: u32,
#[serde(default = "default_budget_run_cap_tokens")]
pub budget_run_cap_tokens: u32,
#[serde(default = "default_true")]
pub ambient: bool,
#[serde(default = "default_true")]
pub compress_capsules: bool,
#[serde(default = "default_true")]
pub session_dedupe: bool,
}
fn default_max_capsules() -> usize {
8
}
fn default_min_semantic_score() -> f32 {
-1.0
}
fn default_min_lexical_coverage() -> f32 {
0.5
}
fn default_budget_floor_tokens() -> u32 {
1500
}
fn default_budget_run_cap_tokens() -> u32 {
8000
}
impl Default for BrokerSection {
fn default() -> Self {
Self {
default_budget_tokens: 6000,
weights: BrokerWeights::default(),
max_capsules: default_max_capsules(),
min_semantic_score: default_min_semantic_score(),
min_lexical_coverage: default_min_lexical_coverage(),
budget_floor_tokens: default_budget_floor_tokens(),
budget_run_cap_tokens: default_budget_run_cap_tokens(),
ambient: default_true(),
compress_capsules: default_true(),
session_dedupe: default_true(),
}
}
}
pub fn adaptive_budget(task_size: u32, floor: u32, run_cap: u32) -> u32 {
if task_size == 0 {
return floor;
}
const K_SCALED: u32 = 3182; let sqrt_part = (task_size as f64).sqrt();
let budget_f = floor as f64 + (K_SCALED as f64 / 10.0) * sqrt_part;
let budget = budget_f.round() as u32;
budget.clamp(floor, run_cap)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrokerWeights {
pub relevance: f32,
pub confidence: f32,
pub freshness: f32,
pub scope: f32,
pub localization: Option<StageWeights>,
pub patch_plan: Option<StageWeights>,
pub verification: Option<StageWeights>,
pub review: Option<StageWeights>,
#[serde(default = "default_decay_half_life_days")]
pub decay_half_life_days: f32,
}
fn default_decay_half_life_days() -> f32 {
30.0
}
impl Default for BrokerWeights {
fn default() -> Self {
Self {
relevance: 0.50,
confidence: 0.20,
freshness: 0.20,
scope: 0.10,
localization: Some(StageWeights {
relevance: 0.70,
confidence: 0.10,
freshness: 0.10,
scope: 0.10,
}),
patch_plan: Some(StageWeights {
relevance: 0.40,
confidence: 0.30,
freshness: 0.10,
scope: 0.20,
}),
verification: Some(StageWeights {
relevance: 0.40,
confidence: 0.10,
freshness: 0.40,
scope: 0.10,
}),
review: Some(StageWeights {
relevance: 0.50,
confidence: 0.20,
freshness: 0.20,
scope: 0.10,
}),
decay_half_life_days: default_decay_half_life_days(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StageWeights {
pub relevance: f32,
pub confidence: f32,
pub freshness: f32,
pub scope: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShellSection {
pub default_timeout_secs: u64,
pub max_timeout_secs: u64,
pub env_allowlist_extra: Vec<String>,
pub redact_secrets: bool,
}
impl Default for ShellSection {
fn default() -> Self {
Self {
default_timeout_secs: 60,
max_timeout_secs: 600,
env_allowlist_extra: vec!["RUSTFLAGS".to_string(), "CARGO_HOME".to_string()],
redact_secrets: true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IngestionSection {
pub max_file_bytes: u64,
pub extra_skip_dirs: Vec<String>,
pub max_total_files: u64,
#[serde(default = "default_true")]
pub detect_conflicts: bool,
}
impl Default for IngestionSection {
fn default() -> Self {
Self {
max_file_bytes: 524_288,
extra_skip_dirs: Vec::new(),
max_total_files: 50_000,
detect_conflicts: true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunSection {
pub max_total_tool_calls: u32,
pub max_total_model_turns: u32,
pub max_total_cost_usd: f32,
}
impl Default for RunSection {
fn default() -> Self {
Self {
max_total_tool_calls: 60,
max_total_model_turns: 30,
max_total_cost_usd: 250.0,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pre_v0_8_config_without_embedder_loads_with_default() {
let toml = r#"
[kimetsu]
project_id = "demo"
schema_version = 7
[model]
provider = "anthropic"
model = "claude-opus-4-7"
api_key_env = "ANTHROPIC_API_KEY"
max_output_tokens = 8192
temperature = 0.2
request_timeout_secs = 120
[broker]
default_budget_tokens = 6000
[broker.weights]
relevance = 0.5
confidence = 0.2
freshness = 0.2
scope = 0.1
[shell]
default_timeout_secs = 60
max_timeout_secs = 600
env_allowlist_extra = []
redact_secrets = true
[ingestion]
max_file_bytes = 524288
extra_skip_dirs = []
max_total_files = 50000
[run]
max_total_tool_calls = 60
max_total_model_turns = 30
max_total_cost_usd = 250.0
"#;
let config = ProjectConfig::from_toml(toml).expect("pre-v0.8 toml must load");
assert_eq!(config.embedder.model, "jina-v2-base-code");
assert!(config.learning.auto_harvest);
assert!(!config.learning.distiller.enabled);
assert_eq!(config.learning.distiller.provider, "anthropic");
assert_eq!(config.learning.distiller.model, "claude-haiku-4-5");
assert_eq!(config.learning.distiller.api_key_env, "ANTHROPIC_API_KEY");
assert_eq!(config.learning.distiller.base_url_env, "ANTHROPIC_BASE_URL");
assert_eq!(config.broker.max_capsules, 8);
assert_eq!(config.broker.min_semantic_score, -1.0, "auto sentinel");
assert_eq!(config.broker.min_lexical_coverage, 0.5);
assert_eq!(config.broker.budget_floor_tokens, 1500);
assert_eq!(config.broker.budget_run_cap_tokens, 8000);
assert!(
config.embedder.enabled,
"W3.1: embedder.enabled must default to true"
);
assert!(
config.broker.ambient,
"W3.2: broker.ambient must default to true"
);
assert!(
config.kimetsu.use_user_brain,
"W3.3: kimetsu.use_user_brain must default to true"
);
assert!(
config.embedder.daemon,
"embedder.daemon must default to true"
);
assert!(
config.embedder.warm_on_start,
"embedder.warm_on_start must default to true"
);
assert!(
config.kimetsu.mcp_write_tools,
"kimetsu.mcp_write_tools must default to true"
);
assert_eq!(
config.embedder.reranker, "ms-marco-tinybert-l-2-v2",
"embedder.reranker must default to ms-marco-tinybert-l-2-v2"
);
assert!(
config.learning.store_queries,
"learning.store_queries must default to true"
);
assert!(
config.broker.compress_capsules,
"broker.compress_capsules must default to true"
);
assert!(
config.broker.session_dedupe,
"broker.session_dedupe must default to true"
);
}
#[test]
fn default_config_uses_config_version_not_schema_version() {
let cfg = ProjectConfig::default_for_project("p1");
assert_eq!(cfg.kimetsu.schema_version, crate::KIMETSU_CONFIG_VERSION);
}
#[test]
fn embedder_survives_toml_round_trip() {
let mut config = ProjectConfig::default_for_project("demo");
config.embedder.model = "bge-m3".to_string();
let serialized = config.to_toml().expect("serialize");
let reloaded = ProjectConfig::from_toml(&serialized).expect("reload");
assert_eq!(reloaded.embedder.model, "bge-m3");
assert_eq!(reloaded.broker.default_budget_tokens, 6000);
assert_eq!(reloaded.kimetsu.project_id, "demo");
assert_eq!(reloaded.broker.budget_floor_tokens, 1500);
assert_eq!(reloaded.broker.budget_run_cap_tokens, 8000);
assert!(reloaded.embedder.enabled);
assert!(reloaded.broker.ambient);
assert!(reloaded.kimetsu.use_user_brain);
}
#[test]
fn w3_off_switch_fields_round_trip_as_false() {
let mut config = ProjectConfig::default_for_project("demo");
config.embedder.enabled = false;
config.broker.ambient = false;
config.kimetsu.use_user_brain = false;
let serialized = config.to_toml().expect("serialize");
let reloaded = ProjectConfig::from_toml(&serialized).expect("reload");
assert!(
!reloaded.embedder.enabled,
"embedder.enabled must survive as false"
);
assert!(
!reloaded.broker.ambient,
"broker.ambient must survive as false"
);
assert!(
!reloaded.kimetsu.use_user_brain,
"kimetsu.use_user_brain must survive as false"
);
assert_eq!(reloaded.kimetsu.project_id, "demo");
}
#[test]
fn f3_adaptive_budget_zero_size_returns_floor() {
assert_eq!(
super::adaptive_budget(0, 1500, 8000),
1500,
"task_size=0 must return floor"
);
}
#[test]
fn f3_adaptive_budget_huge_size_clamped_to_run_cap() {
let result = super::adaptive_budget(1_000_000, 1500, 8000);
assert_eq!(result, 8000, "huge task_size must be clamped to run_cap");
}
#[test]
fn f3_adaptive_budget_is_sublinear() {
let floor = 1500u32;
let run_cap = 16_000u32;
let t0 = 200u32;
let b_t0 = super::adaptive_budget(t0, floor, run_cap);
let b_5t0 = super::adaptive_budget(5 * t0, floor, run_cap);
assert!(
b_5t0 < 2 * b_t0,
"sublinear guarantee: adaptive_budget(5*T)={b_5t0} must be < 2*adaptive_budget(T)={} (T={t0})",
2 * b_t0
);
assert!(
b_5t0 > b_t0,
"budget must still grow: adaptive_budget(5*T)={b_5t0} > adaptive_budget(T)={b_t0}"
);
}
#[test]
fn f3_adaptive_budget_typical_task_near_historical_default() {
let budget = super::adaptive_budget(200, 1500, 8000);
assert!(
(5700..=8000).contains(&budget),
"typical task budget expected near 6000, got {budget}"
);
}
#[test]
fn f3_adaptive_budget_respects_floor() {
for size in [1u32, 5, 10, 50] {
let b = super::adaptive_budget(size, 1500, 8000);
assert!(
b >= 1500,
"task_size={size}: budget={b} must be >= floor=1500"
);
}
}
#[test]
fn f3_adaptive_budget_respects_run_cap() {
for size in [500u32, 1000, 5000, 100_000] {
let b = super::adaptive_budget(size, 1500, 8000);
assert!(
b <= 8000,
"task_size={size}: budget={b} must be <= run_cap=8000"
);
}
}
#[test]
fn pre_v1_5_config_without_price_per_mtok_loads_with_none() {
let toml = r#"
[kimetsu]
project_id = "demo"
schema_version = 7
[model]
provider = "anthropic"
model = "claude-sonnet-4-7"
api_key_env = "ANTHROPIC_API_KEY"
max_output_tokens = 8192
temperature = 0.2
request_timeout_secs = 120
[broker]
default_budget_tokens = 6000
[broker.weights]
relevance = 0.5
confidence = 0.2
freshness = 0.2
scope = 0.1
[shell]
default_timeout_secs = 60
max_timeout_secs = 600
env_allowlist_extra = []
redact_secrets = true
[ingestion]
max_file_bytes = 524288
extra_skip_dirs = []
max_total_files = 50000
[run]
max_total_tool_calls = 60
max_total_model_turns = 30
max_total_cost_usd = 250.0
"#;
let config = ProjectConfig::from_toml(toml).expect("pre-v1.5 toml must load");
assert!(
config.model.price_per_mtok.is_none(),
"price_per_mtok must default to None when absent from project.toml"
);
}
#[test]
fn broker_v1_5_fields_round_trip_as_false() {
let mut config = ProjectConfig::default_for_project("demo");
config.broker.compress_capsules = false;
config.broker.session_dedupe = false;
let serialized = config.to_toml().expect("serialize");
let reloaded = ProjectConfig::from_toml(&serialized).expect("reload");
assert!(
!reloaded.broker.compress_capsules,
"compress_capsules must survive as false"
);
assert!(
!reloaded.broker.session_dedupe,
"session_dedupe must survive as false"
);
}
#[test]
fn price_per_mtok_round_trips() {
let mut config = ProjectConfig::default_for_project("demo");
config.model.price_per_mtok = Some(7.5);
let serialized = config.to_toml().expect("serialize");
let reloaded = ProjectConfig::from_toml(&serialized).expect("reload");
assert_eq!(
reloaded.model.price_per_mtok,
Some(7.5),
"price_per_mtok must round-trip"
);
}
}