storkit: merge 357_story_bot_assign_command_to_pre_assign_a_model_to_a_story
This commit is contained in:
Dave
385
server/src/matrix/commands/assign.rs
Normal file
385
server/src/matrix/commands/assign.rs
Normal file
@@ -0,0 +1,385 @@
|
||||
//! Handler for the `assign` command.
|
||||
//!
|
||||
//! `assign <number> <model>` 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-<model>` 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<String> {
|
||||
let args = ctx.args.trim();
|
||||
|
||||
// Parse `<number> <model>` 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 <number> <model>` (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 <number> <model>`",
|
||||
ctx.bot_name
|
||||
));
|
||||
}
|
||||
|
||||
if model_str.is_empty() {
|
||||
return Some(format!(
|
||||
"Usage: `{} assign <number> <model>` (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: <agent_name>` 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<String> {
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -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 <number> <model>` (e.g. `assign 42 opus`)",
|
||||
handler: assign::handle_assign,
|
||||
},
|
||||
BotCommand {
|
||||
name: "help",
|
||||
description: "Show this list of available commands",
|
||||
|
||||
Reference in New Issue
Block a user