From 83879cfa9efcaf20fff74e9825fb9837a998fd83 Mon Sep 17 00:00:00 2001 From: Dave Date: Fri, 20 Mar 2026 18:49:20 +0000 Subject: [PATCH] storkit: merge 357_story_bot_assign_command_to_pre_assign_a_model_to_a_story --- server/src/matrix/commands/assign.rs | 385 +++++++++++++++++++++++++++ server/src/matrix/commands/mod.rs | 6 + 2 files changed, 391 insertions(+) create mode 100644 server/src/matrix/commands/assign.rs diff --git a/server/src/matrix/commands/assign.rs b/server/src/matrix/commands/assign.rs new file mode 100644 index 0000000..e734733 --- /dev/null +++ b/server/src/matrix/commands/assign.rs @@ -0,0 +1,385 @@ +//! Handler for the `assign` command. +//! +//! `assign ` pre-assigns a coder model (e.g. `opus`, `sonnet`) +//! to a story before it starts. The assignment persists in the story file's +//! front matter as `agent: coder-` so that when the pipeline picks up +//! the story — either via auto-assign or the `start` command — it uses the +//! assigned model instead of the default. + +use super::CommandContext; +use crate::io::story_metadata::{parse_front_matter, set_front_matter_field}; + +/// All pipeline stage directories to search when finding a work item by number. +const STAGES: &[&str] = &[ + "1_backlog", + "2_current", + "3_qa", + "4_merge", + "5_done", + "6_archived", +]; + +/// Resolve a model name hint (e.g. `"opus"`) to a full agent name +/// (e.g. `"coder-opus"`). If the hint already starts with `"coder-"`, +/// it is returned unchanged to prevent double-prefixing. +fn resolve_agent_name(model: &str) -> String { + if model.starts_with("coder-") { + model.to_string() + } else { + format!("coder-{model}") + } +} + +pub(super) fn handle_assign(ctx: &CommandContext) -> Option { + let args = ctx.args.trim(); + + // Parse ` ` from args. + let (number_str, model_str) = match args.split_once(char::is_whitespace) { + Some((n, m)) => (n.trim(), m.trim()), + None => { + return Some(format!( + "Usage: `{} assign ` (e.g. `assign 42 opus`)", + ctx.bot_name + )); + } + }; + + if number_str.is_empty() || !number_str.chars().all(|c| c.is_ascii_digit()) { + return Some(format!( + "Invalid story number `{number_str}`. Usage: `{} assign `", + ctx.bot_name + )); + } + + if model_str.is_empty() { + return Some(format!( + "Usage: `{} assign ` (e.g. `assign 42 opus`)", + ctx.bot_name + )); + } + + // Find the story file across all pipeline stages. + let mut found: Option<(std::path::PathBuf, String)> = None; + 'outer: for stage in STAGES { + let dir = ctx.project_root.join(".storkit").join("work").join(stage); + if !dir.exists() { + continue; + } + if let Ok(entries) = std::fs::read_dir(&dir) { + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().and_then(|e| e.to_str()) != Some("md") { + continue; + } + if let Some(stem) = path + .file_stem() + .and_then(|s| s.to_str()) + .map(|s| s.to_string()) + { + let file_num = stem + .split('_') + .next() + .filter(|s| !s.is_empty() && s.chars().all(|c| c.is_ascii_digit())) + .unwrap_or("") + .to_string(); + if file_num == number_str { + found = Some((path, stem)); + break 'outer; + } + } + } + } + } + + let (path, story_id) = match found { + Some(f) => f, + None => { + return Some(format!( + "No story, bug, or spike with number **{number_str}** found." + )); + } + }; + + // Read the human-readable name from front matter for the response. + let story_name = std::fs::read_to_string(&path) + .ok() + .and_then(|contents| { + parse_front_matter(&contents) + .ok() + .and_then(|m| m.name) + }) + .unwrap_or_else(|| story_id.clone()); + + let agent_name = resolve_agent_name(model_str); + + // Write `agent: ` into the story's front matter. + let result = std::fs::read_to_string(&path) + .map_err(|e| format!("Failed to read story file: {e}")) + .and_then(|contents| { + let updated = set_front_matter_field(&contents, "agent", &agent_name); + std::fs::write(&path, &updated) + .map_err(|e| format!("Failed to write story file: {e}")) + }); + + match result { + Ok(()) => Some(format!( + "Assigned **{agent_name}** to **{story_name}** (story {number_str}). \ + The model will be used when the story starts." + )), + Err(e) => Some(format!( + "Failed to assign model to **{story_name}**: {e}" + )), + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use crate::agents::AgentPool; + use std::collections::HashSet; + use std::sync::{Arc, Mutex}; + + use super::super::{CommandDispatch, try_handle_command}; + + fn assign_cmd_with_root(root: &std::path::Path, args: &str) -> Option { + let agents = Arc::new(AgentPool::new_test(3000)); + let ambient_rooms = Arc::new(Mutex::new(HashSet::new())); + let room_id = "!test:example.com".to_string(); + let dispatch = CommandDispatch { + bot_name: "Timmy", + bot_user_id: "@timmy:homeserver.local", + project_root: root, + agents: &agents, + ambient_rooms: &ambient_rooms, + room_id: &room_id, + }; + try_handle_command(&dispatch, &format!("@timmy assign {args}")) + } + + fn write_story_file(root: &std::path::Path, stage: &str, filename: &str, content: &str) { + let dir = root.join(".storkit/work").join(stage); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write(dir.join(filename), content).unwrap(); + } + + // -- registration / help ------------------------------------------------ + + #[test] + fn assign_command_is_registered() { + use super::super::commands; + let found = commands().iter().any(|c| c.name == "assign"); + assert!(found, "assign command must be in the registry"); + } + + #[test] + fn assign_command_appears_in_help() { + let result = super::super::tests::try_cmd_addressed( + "Timmy", + "@timmy:homeserver.local", + "@timmy help", + ); + let output = result.unwrap(); + assert!( + output.contains("assign"), + "help should list assign command: {output}" + ); + } + + // -- argument validation ------------------------------------------------ + + #[test] + fn assign_no_args_returns_usage() { + let tmp = tempfile::TempDir::new().unwrap(); + let output = assign_cmd_with_root(tmp.path(), "").unwrap(); + assert!( + output.contains("Usage"), + "no args should show usage: {output}" + ); + } + + #[test] + fn assign_missing_model_returns_usage() { + let tmp = tempfile::TempDir::new().unwrap(); + let output = assign_cmd_with_root(tmp.path(), "42").unwrap(); + assert!( + output.contains("Usage"), + "missing model should show usage: {output}" + ); + } + + #[test] + fn assign_non_numeric_number_returns_error() { + let tmp = tempfile::TempDir::new().unwrap(); + let output = assign_cmd_with_root(tmp.path(), "abc opus").unwrap(); + assert!( + output.contains("Invalid story number"), + "non-numeric number should return error: {output}" + ); + } + + // -- story not found ---------------------------------------------------- + + #[test] + fn assign_unknown_story_returns_friendly_message() { + let tmp = tempfile::TempDir::new().unwrap(); + // Create stage dirs but no matching story. + for stage in &["1_backlog", "2_current"] { + std::fs::create_dir_all(tmp.path().join(".storkit/work").join(stage)).unwrap(); + } + let output = assign_cmd_with_root(tmp.path(), "999 opus").unwrap(); + assert!( + output.contains("999") && output.contains("found"), + "not-found message should include number and 'found': {output}" + ); + } + + // -- successful assignment ---------------------------------------------- + + #[test] + fn assign_writes_agent_field_to_front_matter() { + let tmp = tempfile::TempDir::new().unwrap(); + write_story_file( + tmp.path(), + "1_backlog", + "42_story_test_feature.md", + "---\nname: Test Feature\n---\n\n# Story 42\n", + ); + + let output = assign_cmd_with_root(tmp.path(), "42 opus").unwrap(); + assert!( + output.contains("coder-opus"), + "confirmation should include resolved agent name: {output}" + ); + assert!( + output.contains("Test Feature"), + "confirmation should include story name: {output}" + ); + + // Verify the file was updated. + let contents = std::fs::read_to_string( + tmp.path() + .join(".storkit/work/1_backlog/42_story_test_feature.md"), + ) + .unwrap(); + assert!( + contents.contains("agent: coder-opus"), + "front matter should contain agent field: {contents}" + ); + } + + #[test] + fn assign_with_sonnet_writes_coder_sonnet() { + let tmp = tempfile::TempDir::new().unwrap(); + write_story_file( + tmp.path(), + "2_current", + "10_story_current.md", + "---\nname: Current Story\n---\n", + ); + + assign_cmd_with_root(tmp.path(), "10 sonnet").unwrap(); + + let contents = std::fs::read_to_string( + tmp.path() + .join(".storkit/work/2_current/10_story_current.md"), + ) + .unwrap(); + assert!( + contents.contains("agent: coder-sonnet"), + "front matter should contain agent: coder-sonnet: {contents}" + ); + } + + #[test] + fn assign_with_already_prefixed_name_does_not_double_prefix() { + let tmp = tempfile::TempDir::new().unwrap(); + write_story_file( + tmp.path(), + "1_backlog", + "7_story_small.md", + "---\nname: Small Story\n---\n", + ); + + let output = assign_cmd_with_root(tmp.path(), "7 coder-opus").unwrap(); + assert!( + output.contains("coder-opus"), + "should not double-prefix: {output}" + ); + assert!( + !output.contains("coder-coder-opus"), + "must not double-prefix: {output}" + ); + + let contents = std::fs::read_to_string( + tmp.path().join(".storkit/work/1_backlog/7_story_small.md"), + ) + .unwrap(); + assert!( + contents.contains("agent: coder-opus"), + "must write coder-opus, not coder-coder-opus: {contents}" + ); + } + + #[test] + fn assign_overwrites_existing_agent_field() { + let tmp = tempfile::TempDir::new().unwrap(); + write_story_file( + tmp.path(), + "1_backlog", + "5_story_existing.md", + "---\nname: Existing\nagent: coder-sonnet\n---\n", + ); + + assign_cmd_with_root(tmp.path(), "5 opus").unwrap(); + + let contents = std::fs::read_to_string( + tmp.path() + .join(".storkit/work/1_backlog/5_story_existing.md"), + ) + .unwrap(); + assert!( + contents.contains("agent: coder-opus"), + "should overwrite old agent with new: {contents}" + ); + assert!( + !contents.contains("coder-sonnet"), + "old agent should no longer appear: {contents}" + ); + } + + #[test] + fn assign_finds_story_in_any_stage() { + let tmp = tempfile::TempDir::new().unwrap(); + // Story is in 3_qa/, not backlog. + write_story_file( + tmp.path(), + "3_qa", + "99_story_in_qa.md", + "---\nname: In QA\n---\n", + ); + + let output = assign_cmd_with_root(tmp.path(), "99 opus").unwrap(); + assert!( + output.contains("coder-opus"), + "should find story in qa stage: {output}" + ); + } + + // -- resolve_agent_name unit tests -------------------------------------- + + #[test] + fn resolve_agent_name_prefixes_bare_model() { + assert_eq!(super::resolve_agent_name("opus"), "coder-opus"); + assert_eq!(super::resolve_agent_name("sonnet"), "coder-sonnet"); + assert_eq!(super::resolve_agent_name("haiku"), "coder-haiku"); + } + + #[test] + fn resolve_agent_name_does_not_double_prefix() { + assert_eq!(super::resolve_agent_name("coder-opus"), "coder-opus"); + assert_eq!(super::resolve_agent_name("coder-sonnet"), "coder-sonnet"); + } +} diff --git a/server/src/matrix/commands/mod.rs b/server/src/matrix/commands/mod.rs index e5af9b9..019a3e0 100644 --- a/server/src/matrix/commands/mod.rs +++ b/server/src/matrix/commands/mod.rs @@ -6,6 +6,7 @@ //! as they are added. mod ambient; +mod assign; mod cost; mod git; mod help; @@ -75,6 +76,11 @@ pub struct CommandContext<'a> { /// Add new commands here — they will automatically appear in `help` output. pub fn commands() -> &'static [BotCommand] { &[ + BotCommand { + name: "assign", + description: "Pre-assign a model to a story: `assign ` (e.g. `assign 42 opus`)", + handler: assign::handle_assign, + }, BotCommand { name: "help", description: "Show this list of available commands",