story-kit: merge 307_story_configurable_coder_pool_size_and_default_model_in_project_toml

This commit is contained in:
Dave
2026-03-19 15:58:32 +00:00
parent 101b365354
commit 429597cbce
4 changed files with 386 additions and 2 deletions

View File

@@ -1568,6 +1568,27 @@ impl AgentPool {
let preferred_agent =
read_story_front_matter_agent(project_root, stage_dir, story_id);
// Check max_coders limit for the Coder stage before agent selection.
// If the pool is full, all remaining items in this stage wait.
if *stage == PipelineStage::Coder
&& let Some(max) = config.max_coders
{
let agents_lock = match self.agents.lock() {
Ok(a) => a,
Err(e) => {
slog_error!("[auto-assign] Failed to lock agents: {e}");
break;
}
};
let active = count_active_agents_for_stage(&config, &agents_lock, stage);
if active >= max {
slog!(
"[auto-assign] Coder pool full ({active}/{max}); remaining items in {stage_dir}/ will wait."
);
break;
}
}
// Outcome: (already_assigned, chosen_agent, preferred_busy, stage_mismatch)
// preferred_busy=true means the story has a specific agent requested but it is
// currently occupied — the story should wait rather than fall back.
@@ -2225,18 +2246,60 @@ fn is_story_assigned_for_stage(
})
}
/// Count active (pending/running) agents for a given pipeline stage.
fn count_active_agents_for_stage(
config: &ProjectConfig,
agents: &HashMap<String, StoryAgent>,
stage: &PipelineStage,
) -> usize {
agents
.values()
.filter(|a| {
matches!(a.status, AgentStatus::Running | AgentStatus::Pending)
&& config
.find_agent(&a.agent_name)
.map(|cfg| agent_config_stage(cfg) == *stage)
.unwrap_or_else(|| pipeline_stage(&a.agent_name) == *stage)
})
.count()
}
/// Find the first configured agent for `stage` that has no active (pending/running) assignment.
/// Returns `None` if all agents for that stage are busy or none are configured.
/// Uses the agent's explicit `stage` config field (preferred) or falls back to name-based detection.
/// Returns `None` if all agents for that stage are busy, none are configured,
/// or the `max_coders` limit has been reached (for the Coder stage).
///
/// For the Coder stage, when `default_coder_model` is set, only considers agents whose
/// model matches the default. This ensures opus-class agents are reserved for explicit
/// front-matter requests.
fn find_free_agent_for_stage<'a>(
config: &'a ProjectConfig,
agents: &HashMap<String, StoryAgent>,
stage: &PipelineStage,
) -> Option<&'a str> {
// Enforce max_coders limit for the Coder stage.
if *stage == PipelineStage::Coder
&& let Some(max) = config.max_coders
{
let active = count_active_agents_for_stage(config, agents, stage);
if active >= max {
return None;
}
}
for agent_config in &config.agent {
if agent_config_stage(agent_config) != *stage {
continue;
}
// When default_coder_model is set, only auto-assign coder agents whose
// model matches. This keeps opus agents reserved for explicit requests.
if *stage == PipelineStage::Coder
&& let Some(ref default_model) = config.default_coder_model
{
let agent_model = agent_config.model.as_deref().unwrap_or("");
if agent_model != default_model {
continue;
}
}
let is_busy = agents.values().any(|a| {
a.agent_name == agent_config.name
&& matches!(a.status, AgentStatus::Running | AgentStatus::Pending)
@@ -5219,4 +5282,197 @@ stage = "qa"
.collect::<Vec<_>>()
);
}
// ── Helper to construct a test StoryAgent ──────────────────────────
fn make_test_story_agent(agent_name: &str, status: AgentStatus) -> StoryAgent {
StoryAgent {
agent_name: agent_name.to_string(),
status,
worktree_info: None,
session_id: None,
tx: broadcast::channel(1).0,
task_handle: None,
event_log: Arc::new(Mutex::new(Vec::new())),
completion: None,
project_root: None,
log_session_id: None,
merge_failure_reported: false,
}
}
// ── find_free_agent_for_stage: default_coder_model filtering ─────────
#[test]
fn find_free_agent_skips_opus_when_default_coder_model_set() {
let config = make_config(
r#"
default_coder_model = "sonnet"
[[agent]]
name = "coder-1"
stage = "coder"
model = "sonnet"
[[agent]]
name = "coder-opus"
stage = "coder"
model = "opus"
"#,
);
let agents = HashMap::new();
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
assert_eq!(free, Some("coder-1"));
}
#[test]
fn find_free_agent_returns_opus_when_no_default_coder_model() {
let config = make_config(
r#"
[[agent]]
name = "coder-opus"
stage = "coder"
model = "opus"
"#,
);
let agents = HashMap::new();
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
assert_eq!(free, Some("coder-opus"));
}
#[test]
fn find_free_agent_returns_none_when_all_sonnet_coders_busy() {
let config = make_config(
r#"
default_coder_model = "sonnet"
[[agent]]
name = "coder-1"
stage = "coder"
model = "sonnet"
[[agent]]
name = "coder-opus"
stage = "coder"
model = "opus"
"#,
);
let mut agents = HashMap::new();
agents.insert(
"story1:coder-1".to_string(),
make_test_story_agent("coder-1", AgentStatus::Running),
);
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
assert_eq!(free, None, "opus agent should not be auto-assigned");
}
// ── find_free_agent_for_stage: max_coders limit ─────────────────────
#[test]
fn find_free_agent_respects_max_coders() {
let config = make_config(
r#"
max_coders = 1
[[agent]]
name = "coder-1"
stage = "coder"
model = "sonnet"
[[agent]]
name = "coder-2"
stage = "coder"
model = "sonnet"
"#,
);
let mut agents = HashMap::new();
agents.insert(
"story1:coder-1".to_string(),
make_test_story_agent("coder-1", AgentStatus::Running),
);
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
assert_eq!(free, None, "max_coders=1 should block second coder");
}
#[test]
fn find_free_agent_allows_within_max_coders() {
let config = make_config(
r#"
max_coders = 2
[[agent]]
name = "coder-1"
stage = "coder"
model = "sonnet"
[[agent]]
name = "coder-2"
stage = "coder"
model = "sonnet"
"#,
);
let mut agents = HashMap::new();
agents.insert(
"story1:coder-1".to_string(),
make_test_story_agent("coder-1", AgentStatus::Running),
);
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
assert_eq!(free, Some("coder-2"));
}
#[test]
fn max_coders_does_not_affect_qa_stage() {
let config = make_config(
r#"
max_coders = 1
[[agent]]
name = "qa"
stage = "qa"
model = "sonnet"
"#,
);
let agents = HashMap::new();
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Qa);
assert_eq!(free, Some("qa"));
}
// ── count_active_agents_for_stage ────────────────────────────────────
#[test]
fn count_active_agents_counts_running_and_pending() {
let config = make_config(
r#"
[[agent]]
name = "coder-1"
stage = "coder"
[[agent]]
name = "coder-2"
stage = "coder"
"#,
);
let mut agents = HashMap::new();
agents.insert(
"s1:coder-1".to_string(),
make_test_story_agent("coder-1", AgentStatus::Running),
);
agents.insert(
"s2:coder-2".to_string(),
make_test_story_agent("coder-2", AgentStatus::Completed),
);
let count = count_active_agents_for_stage(&config, &agents, &PipelineStage::Coder);
assert_eq!(count, 1, "Only Running coder should be counted, not Completed");
}
}

View File

@@ -15,6 +15,17 @@ pub struct ProjectConfig {
/// Per-story `qa` front matter overrides this. Default: "server".
#[serde(default = "default_qa")]
pub default_qa: String,
/// Default model for coder-stage agents (e.g. "sonnet").
/// When set, `find_free_agent_for_stage` only considers coder agents whose
/// model matches this value, so opus agents are only used when explicitly
/// requested via story front matter `agent:` field.
#[serde(default)]
pub default_coder_model: Option<String>,
/// Maximum number of concurrent coder-stage agents.
/// When set, `auto_assign_available_work` will not start more than this many
/// coder agents at once. Stories wait in `2_current/` until a slot frees up.
#[serde(default)]
pub max_coders: Option<usize>,
}
/// Configuration for the filesystem watcher's sweep behaviour.
@@ -134,6 +145,10 @@ struct LegacyProjectConfig {
watcher: WatcherConfig,
#[serde(default = "default_qa")]
default_qa: String,
#[serde(default)]
default_coder_model: Option<String>,
#[serde(default)]
max_coders: Option<usize>,
}
impl Default for ProjectConfig {
@@ -156,6 +171,8 @@ impl Default for ProjectConfig {
}],
watcher: WatcherConfig::default(),
default_qa: default_qa(),
default_coder_model: None,
max_coders: None,
}
}
}
@@ -198,6 +215,8 @@ impl ProjectConfig {
agent: vec![agent],
watcher: legacy.watcher,
default_qa: legacy.default_qa,
default_coder_model: legacy.default_coder_model,
max_coders: legacy.max_coders,
};
validate_agents(&config.agent)?;
return Ok(config);
@@ -219,6 +238,8 @@ impl ProjectConfig {
agent: vec![agent],
watcher: legacy.watcher,
default_qa: legacy.default_qa,
default_coder_model: legacy.default_coder_model,
max_coders: legacy.max_coders,
};
validate_agents(&config.agent)?;
Ok(config)
@@ -228,6 +249,8 @@ impl ProjectConfig {
agent: Vec::new(),
watcher: legacy.watcher,
default_qa: legacy.default_qa,
default_coder_model: legacy.default_coder_model,
max_coders: legacy.max_coders,
})
}
}
@@ -710,4 +733,68 @@ command = "claude"
assert_eq!(config.watcher.done_retention_secs, 900);
assert_eq!(config.agent.len(), 1);
}
// ── default_coder_model & max_coders ─────────────────────────────────
#[test]
fn parse_default_coder_model_and_max_coders() {
let toml_str = r#"
default_coder_model = "sonnet"
max_coders = 3
[[agent]]
name = "coder-1"
stage = "coder"
model = "sonnet"
[[agent]]
name = "coder-opus"
stage = "coder"
model = "opus"
"#;
let config = ProjectConfig::parse(toml_str).unwrap();
assert_eq!(config.default_coder_model, Some("sonnet".to_string()));
assert_eq!(config.max_coders, Some(3));
}
#[test]
fn default_coder_model_and_max_coders_default_to_none() {
let toml_str = r#"
[[agent]]
name = "coder-1"
"#;
let config = ProjectConfig::parse(toml_str).unwrap();
assert_eq!(config.default_coder_model, None);
assert_eq!(config.max_coders, None);
}
#[test]
fn project_toml_has_default_coder_model_and_max_coders() {
// Verify the actual project.toml has the new settings.
let manifest_dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR"));
let project_root = manifest_dir.parent().unwrap();
let config = ProjectConfig::load(project_root).unwrap();
assert_eq!(config.default_coder_model, Some("sonnet".to_string()));
assert_eq!(config.max_coders, Some(3));
}
#[test]
fn project_toml_has_three_sonnet_coders() {
let manifest_dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR"));
let project_root = manifest_dir.parent().unwrap();
let config = ProjectConfig::load(project_root).unwrap();
let sonnet_coders: Vec<_> = config
.agent
.iter()
.filter(|a| a.stage.as_deref() == Some("coder") && a.model.as_deref() == Some("sonnet"))
.collect();
assert_eq!(
sonnet_coders.len(),
3,
"Expected 3 sonnet coders (coder-1, coder-2, coder-3), found {}",
sonnet_coders.len()
);
}
}

View File

@@ -508,6 +508,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
// Should complete without panic
run_setup_commands(tmp.path(), &config).await;
@@ -526,6 +528,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
// Should complete without panic
run_setup_commands(tmp.path(), &config).await;
@@ -544,6 +548,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
// Setup command failures are non-fatal — should not panic or propagate
run_setup_commands(tmp.path(), &config).await;
@@ -562,6 +568,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
// Teardown failures are best-effort — should not propagate
assert!(run_teardown_commands(tmp.path(), &config).await.is_ok());
@@ -579,6 +587,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
let info = create_worktree(&project_root, "42_fresh_test", &config, 3001)
.await
@@ -603,6 +613,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
// First creation
let _info1 = create_worktree(&project_root, "43_reuse_test", &config, 3001)
@@ -643,6 +655,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
let result = remove_worktree_by_story_id(tmp.path(), "99_nonexistent", &config).await;
@@ -666,6 +680,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
create_worktree(&project_root, "88_remove_by_id", &config, 3001)
.await
@@ -720,6 +736,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
// Even though setup commands fail, create_worktree must succeed
// so the agent can start and fix the problem itself.
@@ -746,6 +764,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
// First creation — no setup commands, should succeed
create_worktree(&project_root, "173_reuse_fail", &empty_config, 3001)
@@ -762,6 +782,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
// Second call — worktree exists, setup commands fail, must still succeed
let result =
@@ -785,6 +807,8 @@ mod tests {
agent: vec![],
watcher: WatcherConfig::default(),
default_qa: "server".to_string(),
default_coder_model: None,
max_coders: None,
};
let info = create_worktree(&project_root, "77_remove_async", &config, 3001)
.await