1153 lines
36 KiB
Rust
1153 lines
36 KiB
Rust
use super::*;
|
|
|
|
use crate::client::{Message, MessageContent, MessageRole, Model};
|
|
use crate::config::builtin_manifest;
|
|
use crate::function::write_file_atomic;
|
|
use crate::hooks::HooksMap;
|
|
|
|
use anyhow::Result;
|
|
use fancy_regex::Regex;
|
|
use rust_embed::Embed;
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::Value;
|
|
use std::collections::BTreeSet;
|
|
use std::path::Path;
|
|
use std::sync::LazyLock;
|
|
|
|
pub const SHELL_ROLE: &str = "shell";
|
|
pub const EXPLAIN_SHELL_ROLE: &str = "explain-shell";
|
|
pub const CODE_ROLE: &str = "code";
|
|
pub const CREATE_TITLE_ROLE: &str = "create-title";
|
|
|
|
pub const INPUT_PLACEHOLDER: &str = "__INPUT__";
|
|
|
|
#[derive(Embed)]
|
|
#[folder = "assets/roles/"]
|
|
struct RolesAsset;
|
|
|
|
static RE_METADATA: LazyLock<Regex> =
|
|
LazyLock::new(|| Regex::new(r"(?s)\A\s*-{3,}\s*(.*?)\s*-{3,}\s*(.*)").unwrap());
|
|
|
|
pub trait RoleLike {
|
|
fn to_role(&self) -> Role;
|
|
fn model(&self) -> &Model;
|
|
fn temperature(&self) -> Option<f64>;
|
|
fn top_p(&self) -> Option<f64>;
|
|
fn enabled_tools(&self) -> Option<Vec<String>>;
|
|
fn enabled_mcp_servers(&self) -> Option<Vec<String>>;
|
|
fn mcp_tools(&self) -> Option<IndexMap<String, Vec<String>>>;
|
|
fn set_model(&mut self, model: Model);
|
|
fn set_temperature(&mut self, value: Option<f64>);
|
|
fn reasoning_effort(&self) -> Option<String>;
|
|
fn set_top_p(&mut self, value: Option<f64>);
|
|
fn set_reasoning_effort(&mut self, value: Option<String>);
|
|
fn set_enabled_tools(&mut self, value: Option<Vec<String>>);
|
|
fn set_enabled_mcp_servers(&mut self, value: Option<Vec<String>>);
|
|
fn set_mcp_tools(&mut self, value: Option<IndexMap<String, Vec<String>>>);
|
|
}
|
|
|
|
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
|
pub struct Role {
|
|
name: String,
|
|
#[serde(default)]
|
|
prompt: String,
|
|
#[serde(
|
|
rename(serialize = "model", deserialize = "model"),
|
|
skip_serializing_if = "Option::is_none"
|
|
)]
|
|
model_id: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
temperature: Option<f64>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
top_p: Option<f64>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
reasoning_effort: Option<String>,
|
|
#[serde(
|
|
default,
|
|
skip_serializing_if = "Option::is_none",
|
|
deserialize_with = "super::deserialize_csv_or_vec"
|
|
)]
|
|
enabled_tools: Option<Vec<String>>,
|
|
#[serde(
|
|
default,
|
|
skip_serializing_if = "Option::is_none",
|
|
deserialize_with = "super::deserialize_csv_or_vec"
|
|
)]
|
|
enabled_mcp_servers: Option<Vec<String>>,
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
skills_enabled: Option<bool>,
|
|
#[serde(
|
|
default,
|
|
skip_serializing_if = "Option::is_none",
|
|
deserialize_with = "super::deserialize_csv_or_vec"
|
|
)]
|
|
enabled_skills: Option<Vec<String>>,
|
|
#[serde(
|
|
default,
|
|
skip_serializing_if = "Option::is_none",
|
|
deserialize_with = "super::deserialize_csv_or_vec"
|
|
)]
|
|
enabled_macros: Option<Vec<String>>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
auto_continue: Option<bool>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
max_auto_continues: Option<usize>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
inject_todo_instructions: Option<bool>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
continuation_prompt: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
inject_skill_instructions: Option<bool>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
skill_instructions: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
memory: Option<bool>,
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
hooks: Option<HooksMap>,
|
|
|
|
#[serde(skip)]
|
|
model: Model,
|
|
}
|
|
|
|
impl Role {
|
|
pub fn new(name: &str, content: &str) -> Self {
|
|
let mut metadata = "";
|
|
let mut prompt = content.trim();
|
|
if let Ok(Some(caps)) = RE_METADATA.captures(content)
|
|
&& let (Some(metadata_value), Some(prompt_value)) = (caps.get(1), caps.get(2))
|
|
{
|
|
let candidate = metadata_value.as_str().trim();
|
|
if candidate.is_empty()
|
|
|| serde_yaml::from_str::<Value>(candidate)
|
|
.map(|v| v.is_object())
|
|
.unwrap_or(false)
|
|
{
|
|
metadata = candidate;
|
|
prompt = prompt_value.as_str().trim();
|
|
}
|
|
}
|
|
let mut prompt = prompt.to_string();
|
|
interpolate_variables(&mut prompt);
|
|
let mut role = Self {
|
|
name: name.to_string(),
|
|
prompt,
|
|
..Default::default()
|
|
};
|
|
if !metadata.is_empty()
|
|
&& let Ok(value) = serde_yaml::from_str::<Value>(metadata)
|
|
&& let Some(value) = value.as_object()
|
|
{
|
|
for (key, value) in value {
|
|
match key.as_str() {
|
|
"model" => role.model_id = value.as_str().map(|v| v.to_string()),
|
|
"temperature" => role.temperature = value.as_f64(),
|
|
"top_p" => role.top_p = value.as_f64(),
|
|
"reasoning_effort" => {
|
|
role.reasoning_effort = value.as_str().map(|v| v.to_string())
|
|
}
|
|
"enabled_tools" => role.enabled_tools = parse_string_or_array(value),
|
|
"enabled_mcp_servers" => {
|
|
role.enabled_mcp_servers = parse_string_or_array(value)
|
|
}
|
|
"mcp_tools" => role.mcp_tools = parse_mcp_tools_map(value),
|
|
"skills_enabled" => role.skills_enabled = value.as_bool(),
|
|
"enabled_skills" => role.enabled_skills = parse_string_or_array(value),
|
|
"enabled_macros" => role.enabled_macros = parse_string_or_array(value),
|
|
"auto_continue" => role.auto_continue = value.as_bool(),
|
|
"max_auto_continues" => {
|
|
role.max_auto_continues = value.as_u64().map(|v| v as usize)
|
|
}
|
|
"inject_todo_instructions" => role.inject_todo_instructions = value.as_bool(),
|
|
"continuation_prompt" => {
|
|
role.continuation_prompt = value.as_str().map(|v| v.to_string())
|
|
}
|
|
"inject_skill_instructions" => role.inject_skill_instructions = value.as_bool(),
|
|
"skill_instructions" => {
|
|
role.skill_instructions = value.as_str().map(|v| v.to_string())
|
|
}
|
|
"memory" => role.memory = value.as_bool(),
|
|
"hooks" => role.hooks = parse_hooks_map(value),
|
|
_ => (),
|
|
}
|
|
}
|
|
}
|
|
role
|
|
}
|
|
|
|
pub fn builtin(name: &str) -> Result<Self> {
|
|
let content = RolesAsset::get(&format!("{name}.md"))
|
|
.ok_or_else(|| anyhow!("Unknown role `{name}`"))?;
|
|
let content = unsafe { std::str::from_utf8_unchecked(&content.data) };
|
|
Ok(Role::new(name, content))
|
|
}
|
|
|
|
pub fn list_builtin_role_names() -> Vec<String> {
|
|
RolesAsset::iter()
|
|
.filter_map(|v| v.strip_suffix(".md").map(|v| v.to_string()))
|
|
.collect()
|
|
}
|
|
|
|
/// Installs builtin role hook scripts into `<roles_dir>/hooks/` and
|
|
/// reconciles ones a release stopped shipping. This is the only builtin
|
|
/// install machinery roles have: general role installation deliberately
|
|
/// does not exist (builtin roles load straight from the embed).
|
|
pub fn install_builtin_role_hooks(force: bool) -> Result<()> {
|
|
let shipped: Vec<(String, String)> = RolesAsset::iter()
|
|
.filter_map(|file| {
|
|
let name = file.as_ref().strip_prefix("hooks/")?;
|
|
if name.is_empty() || name.contains('/') {
|
|
return None;
|
|
}
|
|
let embedded = RolesAsset::get(file.as_ref())?;
|
|
let content = unsafe { std::str::from_utf8_unchecked(&embedded.data) };
|
|
Some((name.to_string(), content.to_string()))
|
|
})
|
|
.collect();
|
|
|
|
install_and_reconcile_role_hooks(&paths::roles_dir().join("hooks"), &shipped, force)
|
|
}
|
|
|
|
pub fn has_args(&self) -> bool {
|
|
self.name.contains('#')
|
|
}
|
|
|
|
pub fn export(&self) -> String {
|
|
let mut metadata = vec![];
|
|
if let Some(model) = self.model_id() {
|
|
metadata.push(format!("model: {model}"));
|
|
}
|
|
if let Some(temperature) = self.temperature() {
|
|
metadata.push(format!("temperature: {temperature}"));
|
|
}
|
|
if let Some(top_p) = self.top_p() {
|
|
metadata.push(format!("top_p: {top_p}"));
|
|
}
|
|
if let Some(reasoning_effort) = self.reasoning_effort() {
|
|
metadata.push(format!("reasoning_effort: {reasoning_effort}"));
|
|
}
|
|
if let Some(enabled_tools) = &self.enabled_tools {
|
|
let inline = serde_json::to_string(enabled_tools).unwrap_or_else(|_| "[]".to_string());
|
|
metadata.push(format!("enabled_tools: {inline}"));
|
|
}
|
|
if let Some(enabled_mcp_servers) = &self.enabled_mcp_servers {
|
|
let inline =
|
|
serde_json::to_string(enabled_mcp_servers).unwrap_or_else(|_| "[]".to_string());
|
|
metadata.push(format!("enabled_mcp_servers: {inline}"));
|
|
}
|
|
if let Some(mcp_tools) = &self.mcp_tools {
|
|
let inline = serde_json::to_string(mcp_tools).unwrap_or_else(|_| "{}".to_string());
|
|
metadata.push(format!("mcp_tools: {inline}"));
|
|
}
|
|
if let Some(skills_enabled) = self.skills_enabled {
|
|
metadata.push(format!("skills_enabled: {skills_enabled}"));
|
|
}
|
|
if let Some(enabled_skills) = &self.enabled_skills {
|
|
let inline = serde_json::to_string(enabled_skills).unwrap_or_else(|_| "[]".to_string());
|
|
metadata.push(format!("enabled_skills: {inline}"));
|
|
}
|
|
if let Some(enabled_macros) = &self.enabled_macros {
|
|
let inline = serde_json::to_string(enabled_macros).unwrap_or_else(|_| "[]".to_string());
|
|
metadata.push(format!("enabled_macros: {inline}"));
|
|
}
|
|
if let Some(auto_continue) = self.auto_continue {
|
|
metadata.push(format!("auto_continue: {auto_continue}"));
|
|
}
|
|
if let Some(max_auto_continues) = self.max_auto_continues {
|
|
metadata.push(format!("max_auto_continues: {max_auto_continues}"));
|
|
}
|
|
if let Some(inject_todo_instructions) = self.inject_todo_instructions {
|
|
metadata.push(format!(
|
|
"inject_todo_instructions: {inject_todo_instructions}"
|
|
));
|
|
}
|
|
if let Some(continuation_prompt) = &self.continuation_prompt {
|
|
metadata.push(format!("continuation_prompt: {continuation_prompt}"));
|
|
}
|
|
if let Some(inject_skill_instructions) = self.inject_skill_instructions {
|
|
metadata.push(format!(
|
|
"inject_skill_instructions: {inject_skill_instructions}"
|
|
));
|
|
}
|
|
if let Some(skill_instructions) = &self.skill_instructions {
|
|
metadata.push(format!("skill_instructions: {skill_instructions}"));
|
|
}
|
|
if let Some(memory) = self.memory {
|
|
metadata.push(format!("memory: {memory}"));
|
|
}
|
|
if let Some(hooks) = &self.hooks {
|
|
let inline = serde_json::to_string(hooks).unwrap_or_else(|_| "{}".to_string());
|
|
metadata.push(format!("hooks: {inline}"));
|
|
}
|
|
if metadata.is_empty() {
|
|
format!("{}\n", self.prompt)
|
|
} else if self.prompt.is_empty() {
|
|
format!("---\n{}\n---\n", metadata.join("\n"))
|
|
} else {
|
|
format!("---\n{}\n---\n\n{}\n", metadata.join("\n"), self.prompt)
|
|
}
|
|
}
|
|
|
|
pub fn save(&mut self, role_name: &str, role_path: &Path, is_repl: bool) -> Result<()> {
|
|
ensure_parent_exists(role_path)?;
|
|
|
|
let content = self.export();
|
|
std::fs::write(role_path, content).with_context(|| {
|
|
format!(
|
|
"Failed to write role {} to {}",
|
|
self.name,
|
|
role_path.display()
|
|
)
|
|
})?;
|
|
|
|
if is_repl {
|
|
println!("✓ Saved role to '{}'.", role_path.display());
|
|
}
|
|
|
|
if role_name != self.name {
|
|
self.name = role_name.to_string();
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn sync<T: RoleLike>(&mut self, role_like: &T) {
|
|
let model = role_like.model();
|
|
let reasoning_effort = role_like.reasoning_effort();
|
|
let temperature = role_like.temperature();
|
|
let top_p = role_like.top_p();
|
|
let enabled_tools = role_like.enabled_tools();
|
|
let enabled_mcp_servers = role_like.enabled_mcp_servers();
|
|
self.batch_set(
|
|
model,
|
|
reasoning_effort,
|
|
temperature,
|
|
top_p,
|
|
enabled_tools,
|
|
enabled_mcp_servers,
|
|
);
|
|
let mcp_tools = role_like.mcp_tools();
|
|
if mcp_tools.is_some() {
|
|
self.set_mcp_tools(mcp_tools);
|
|
}
|
|
}
|
|
|
|
pub fn batch_set(
|
|
&mut self,
|
|
model: &Model,
|
|
reasoning_effort: Option<String>,
|
|
temperature: Option<f64>,
|
|
top_p: Option<f64>,
|
|
enabled_tools: Option<Vec<String>>,
|
|
enabled_mcp_servers: Option<Vec<String>>,
|
|
) {
|
|
self.set_model(model.clone());
|
|
if reasoning_effort.is_some() {
|
|
self.set_reasoning_effort(reasoning_effort.clone());
|
|
}
|
|
if temperature.is_some() {
|
|
self.set_temperature(temperature);
|
|
}
|
|
if top_p.is_some() {
|
|
self.set_top_p(top_p);
|
|
}
|
|
if enabled_tools.is_some() {
|
|
self.set_enabled_tools(enabled_tools);
|
|
}
|
|
if enabled_mcp_servers.is_some() {
|
|
self.set_enabled_mcp_servers(enabled_mcp_servers);
|
|
}
|
|
}
|
|
|
|
pub fn is_derived(&self) -> bool {
|
|
self.name.is_empty()
|
|
}
|
|
|
|
pub fn name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
pub fn model_id(&self) -> Option<&str> {
|
|
self.model_id.as_deref()
|
|
}
|
|
|
|
pub fn prompt(&self) -> &str {
|
|
&self.prompt
|
|
}
|
|
|
|
pub fn is_empty_prompt(&self) -> bool {
|
|
self.prompt.is_empty()
|
|
}
|
|
|
|
pub fn is_embedded_prompt(&self) -> bool {
|
|
self.prompt.contains(INPUT_PLACEHOLDER)
|
|
}
|
|
|
|
pub fn auto_continue(&self) -> Option<bool> {
|
|
self.auto_continue
|
|
}
|
|
|
|
pub fn max_auto_continues(&self) -> Option<usize> {
|
|
self.max_auto_continues
|
|
}
|
|
|
|
pub fn inject_todo_instructions(&self) -> Option<bool> {
|
|
self.inject_todo_instructions
|
|
}
|
|
|
|
pub fn continuation_prompt(&self) -> Option<&str> {
|
|
self.continuation_prompt.as_deref()
|
|
}
|
|
|
|
pub fn inject_skill_instructions(&self) -> Option<bool> {
|
|
self.inject_skill_instructions
|
|
}
|
|
|
|
pub fn skill_instructions(&self) -> Option<&str> {
|
|
self.skill_instructions.as_deref()
|
|
}
|
|
|
|
pub fn memory(&self) -> Option<bool> {
|
|
self.memory
|
|
}
|
|
|
|
pub fn hooks(&self) -> Option<&HooksMap> {
|
|
self.hooks.as_ref()
|
|
}
|
|
|
|
pub fn skills_enabled(&self) -> Option<bool> {
|
|
self.skills_enabled
|
|
}
|
|
|
|
pub fn enabled_skills(&self) -> Option<&[String]> {
|
|
self.enabled_skills.as_deref()
|
|
}
|
|
|
|
pub fn enabled_macros(&self) -> Option<&[String]> {
|
|
self.enabled_macros.as_deref()
|
|
}
|
|
|
|
pub fn append_to_prompt(&mut self, text: &str) {
|
|
self.prompt.push_str(text);
|
|
}
|
|
|
|
pub fn echo_messages(&self, input: &Input) -> String {
|
|
let input_markdown = input.render();
|
|
if self.is_empty_prompt() {
|
|
input_markdown
|
|
} else if self.is_embedded_prompt() {
|
|
self.prompt.replace(INPUT_PLACEHOLDER, &input_markdown)
|
|
} else {
|
|
format!("{}\n\n{}", self.prompt, input_markdown)
|
|
}
|
|
}
|
|
|
|
pub fn build_messages(&self, input: &Input) -> Vec<Message> {
|
|
let mut content = input.message_content();
|
|
let mut messages = if self.is_empty_prompt() {
|
|
vec![Message::new(MessageRole::User, content)]
|
|
} else if self.is_embedded_prompt() {
|
|
content.merge_prompt(|v: &str| self.prompt.replace(INPUT_PLACEHOLDER, v));
|
|
vec![Message::new(MessageRole::User, content)]
|
|
} else {
|
|
let mut messages = vec![];
|
|
let (system, cases) = parse_structure_prompt(&self.prompt);
|
|
if !system.is_empty() {
|
|
messages.push(Message::new(
|
|
MessageRole::System,
|
|
MessageContent::Text(system.to_string()),
|
|
));
|
|
}
|
|
if !cases.is_empty() {
|
|
messages.extend(cases.into_iter().flat_map(|(i, o)| {
|
|
vec![
|
|
Message::new(MessageRole::User, MessageContent::Text(i.to_string())),
|
|
Message::new(MessageRole::Assistant, MessageContent::Text(o.to_string())),
|
|
]
|
|
}));
|
|
}
|
|
messages.push(Message::new(MessageRole::User, content));
|
|
messages
|
|
};
|
|
if let Some(text) = input.continue_output() {
|
|
messages.push(Message::new(
|
|
MessageRole::Assistant,
|
|
MessageContent::Text(text.into()),
|
|
));
|
|
}
|
|
messages
|
|
}
|
|
}
|
|
|
|
impl RoleLike for Role {
|
|
fn to_role(&self) -> Role {
|
|
self.clone()
|
|
}
|
|
|
|
fn model(&self) -> &Model {
|
|
&self.model
|
|
}
|
|
|
|
fn temperature(&self) -> Option<f64> {
|
|
self.temperature
|
|
}
|
|
|
|
fn top_p(&self) -> Option<f64> {
|
|
self.top_p
|
|
}
|
|
|
|
fn reasoning_effort(&self) -> Option<String> {
|
|
self.reasoning_effort.clone()
|
|
}
|
|
|
|
fn enabled_tools(&self) -> Option<Vec<String>> {
|
|
self.enabled_tools.clone()
|
|
}
|
|
|
|
fn enabled_mcp_servers(&self) -> Option<Vec<String>> {
|
|
self.enabled_mcp_servers.clone()
|
|
}
|
|
|
|
fn mcp_tools(&self) -> Option<IndexMap<String, Vec<String>>> {
|
|
self.mcp_tools.clone()
|
|
}
|
|
|
|
fn set_model(&mut self, model: Model) {
|
|
if !self.model().id().is_empty() {
|
|
self.model_id = Some(model.id().to_string());
|
|
}
|
|
self.model = model;
|
|
}
|
|
|
|
fn set_temperature(&mut self, value: Option<f64>) {
|
|
self.temperature = value;
|
|
}
|
|
|
|
fn set_top_p(&mut self, value: Option<f64>) {
|
|
self.top_p = value;
|
|
}
|
|
|
|
fn set_reasoning_effort(&mut self, value: Option<String>) {
|
|
self.reasoning_effort = value;
|
|
}
|
|
|
|
fn set_enabled_tools(&mut self, value: Option<Vec<String>>) {
|
|
self.enabled_tools = value;
|
|
}
|
|
|
|
fn set_enabled_mcp_servers(&mut self, value: Option<Vec<String>>) {
|
|
self.enabled_mcp_servers = value;
|
|
}
|
|
|
|
fn set_mcp_tools(&mut self, value: Option<IndexMap<String, Vec<String>>>) {
|
|
self.mcp_tools = value;
|
|
}
|
|
}
|
|
|
|
fn parse_string_or_array(value: &Value) -> Option<Vec<String>> {
|
|
if value.is_null() {
|
|
return None;
|
|
}
|
|
|
|
if let Some(s) = value.as_str() {
|
|
return Some(csv_to_vec(s));
|
|
}
|
|
|
|
if let Some(arr) = value.as_array() {
|
|
let items: Vec<String> = arr
|
|
.iter()
|
|
.filter_map(|v| v.as_str().map(|s| s.trim().to_string()))
|
|
.filter(|s| !s.is_empty())
|
|
.collect();
|
|
|
|
return Some(items);
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
fn parse_mcp_tools_map(value: &Value) -> Option<IndexMap<String, Vec<String>>> {
|
|
let map = value.as_object()?;
|
|
let mut mcp_tools = IndexMap::new();
|
|
for (server, tools) in map {
|
|
if tools.is_null() {
|
|
mcp_tools.insert(server.clone(), Vec::new());
|
|
} else if let Some(tools) = parse_string_or_array(tools) {
|
|
mcp_tools.insert(server.clone(), tools);
|
|
}
|
|
}
|
|
Some(mcp_tools)
|
|
}
|
|
|
|
fn parse_hooks_map(value: &Value) -> Option<HooksMap> {
|
|
match serde_json::from_value(value.clone()) {
|
|
Ok(map) => Some(map),
|
|
Err(err) => {
|
|
debug!("Ignoring malformed hooks configuration: {err}");
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
fn parse_structure_prompt(prompt: &str) -> (&str, Vec<(&str, &str)>) {
|
|
let mut text = prompt;
|
|
let mut search_input = true;
|
|
let mut system = None;
|
|
let mut parts = vec![];
|
|
loop {
|
|
let search = if search_input {
|
|
"### INPUT:"
|
|
} else {
|
|
"### OUTPUT:"
|
|
};
|
|
match text.find(search) {
|
|
Some(idx) => {
|
|
if system.is_none() {
|
|
system = Some(&text[..idx])
|
|
} else {
|
|
parts.push(&text[..idx])
|
|
}
|
|
search_input = !search_input;
|
|
text = &text[(idx + search.len())..];
|
|
}
|
|
None => {
|
|
if !text.is_empty() {
|
|
if system.is_none() {
|
|
system = Some(text)
|
|
} else {
|
|
parts.push(text)
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
let parts_len = parts.len();
|
|
if parts_len > 0 && parts_len % 2 == 0 {
|
|
let cases: Vec<(&str, &str)> = parts
|
|
.iter()
|
|
.step_by(2)
|
|
.zip(parts.iter().skip(1).step_by(2))
|
|
.map(|(i, o)| (i.trim(), o.trim()))
|
|
.collect();
|
|
let system = system.map(|v| v.trim()).unwrap_or_default();
|
|
return (system, cases);
|
|
}
|
|
|
|
(prompt, vec![])
|
|
}
|
|
|
|
/// Only direct children of the hooks/ directory are supported: they are the
|
|
/// only shape the builtin manifest tracks, so nested asset paths are skipped
|
|
/// by the caller rather than installed unreconciled.
|
|
fn install_and_reconcile_role_hooks(
|
|
dir: &Path,
|
|
shipped: &[(String, String)],
|
|
force: bool,
|
|
) -> Result<()> {
|
|
if !shipped.is_empty() {
|
|
info!("Installing built-in role hooks in {}", dir.display());
|
|
}
|
|
let mut written = BTreeSet::new();
|
|
for (name, content) in shipped {
|
|
let path = dir.join(name);
|
|
if path.exists() && !force {
|
|
debug!(
|
|
"Role hook file already exists, skipping: {}",
|
|
path.display()
|
|
);
|
|
continue;
|
|
}
|
|
ensure_parent_exists(&path)?;
|
|
info!("Creating role hook file: {}", path.display());
|
|
write_file_atomic(&path, content, Some(0o755))?;
|
|
written.insert(name.clone());
|
|
}
|
|
|
|
let names: BTreeSet<String> = shipped.iter().map(|(name, _)| name.clone()).collect();
|
|
// Reconciliation is best-effort housekeeping: a failure here must not
|
|
// abort startup (install_builtins), matching the agent-side policy.
|
|
if let Err(err) = builtin_manifest::reconcile_builtin_dir(dir, &names, &written) {
|
|
warn!(
|
|
"Failed to reconcile builtin role hooks in {}: {err}",
|
|
dir.display()
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn role_new_parses_prompt() {
|
|
let role = Role::new("test", "You are a helpful assistant");
|
|
|
|
assert_eq!(role.name(), "test");
|
|
assert_eq!(role.prompt(), "You are a helpful assistant");
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_parses_metadata() {
|
|
let content =
|
|
"---\nmodel: openai:gpt-4\ntemperature: 0.7\ntop_p: 0.9\n---\nYou are helpful";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
assert_eq!(role.model_id(), Some("openai:gpt-4"));
|
|
assert_eq!(role.temperature(), Some(0.7));
|
|
assert_eq!(role.top_p(), Some(0.9));
|
|
assert_eq!(role.prompt(), "You are helpful");
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_parses_enabled_tools() {
|
|
let content = "---\nenabled_tools: tool1,tool2\n---\nPrompt";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
assert_eq!(
|
|
role.enabled_tools(),
|
|
Some(vec!["tool1".to_string(), "tool2".to_string()])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_parses_enabled_mcp_servers() {
|
|
let content = "---\nenabled_mcp_servers: github,jira\n---\nPrompt";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
assert_eq!(
|
|
role.enabled_mcp_servers(),
|
|
Some(vec!["github".to_string(), "jira".to_string()])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn role_hooks_round_trip_through_export() {
|
|
let content = "---\nhooks:\n turn.completed:\n - name: notify\n command: ./hooks/notify.sh\n---\nPrompt";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
let hooks = role.hooks().expect("hooks should be parsed");
|
|
assert_eq!(hooks["turn.completed"][0].name, "notify");
|
|
assert_eq!(hooks["turn.completed"][0].command, "./hooks/notify.sh");
|
|
|
|
let exported = role.export();
|
|
assert!(exported.contains("hooks:"));
|
|
|
|
let reparsed = Role::new("test", &exported);
|
|
assert_eq!(reparsed.hooks(), role.hooks());
|
|
assert_eq!(reparsed.prompt(), "Prompt");
|
|
}
|
|
|
|
#[test]
|
|
fn role_without_hooks_has_none_and_export_omits_key() {
|
|
let role = Role::new("test", "---\nmemory: true\n---\nPrompt");
|
|
|
|
assert!(role.hooks().is_none());
|
|
assert!(!role.export().contains("hooks:"));
|
|
}
|
|
|
|
#[test]
|
|
fn role_malformed_hooks_degrades_to_none() {
|
|
let content = "---\nhooks: [not, a, map]\n---\nPrompt";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
assert!(role.hooks().is_none());
|
|
assert_eq!(role.prompt(), "Prompt");
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_no_metadata_has_none_fields() {
|
|
let role = Role::new("test", "Just a prompt");
|
|
|
|
assert_eq!(role.model_id(), None);
|
|
assert_eq!(role.temperature(), None);
|
|
assert_eq!(role.top_p(), None);
|
|
assert_eq!(role.enabled_tools(), None);
|
|
assert_eq!(role.enabled_mcp_servers(), None);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_keeps_prompt_containing_markdown_table_separator() {
|
|
let content = "You review PRs.\nThe PR: <url>\n\n| Lane | Verdict |\n|------|---------|\n| code-reviewer | PASS |";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
assert_eq!(role.prompt(), content);
|
|
assert_eq!(role.model_id(), None);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_parses_metadata_with_table_separator_in_prompt() {
|
|
let content =
|
|
"---\nmodel: openai:gpt-4\n---\nIntro kept\n\n| a | b |\n|---|---|\n| 1 | 2 |";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
assert_eq!(role.model_id(), Some("openai:gpt-4"));
|
|
assert_eq!(
|
|
role.prompt(),
|
|
"Intro kept\n\n| a | b |\n|---|---|\n| 1 | 2 |"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_leading_dash_decoration_is_not_metadata() {
|
|
let content = "--- IMPORTANT ---\nDo the thing";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
assert_eq!(role.prompt(), content);
|
|
assert_eq!(role.model_id(), None);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_enabled_macros_absent_is_none() {
|
|
let role = Role::new("test", "---\ntemperature: 0.5\n---\nPrompt");
|
|
|
|
assert_eq!(role.enabled_macros, None);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_enabled_macros_empty_string_is_some_empty() {
|
|
let role = Role::new("test", "---\nenabled_macros: \"\"\n---\nPrompt");
|
|
|
|
assert_eq!(role.enabled_macros, Some(vec![]));
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_enabled_macros_csv_string() {
|
|
let role = Role::new("test", "---\nenabled_macros: a, b\n---\nPrompt");
|
|
|
|
assert_eq!(
|
|
role.enabled_macros,
|
|
Some(vec!["a".to_string(), "b".to_string()])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_enabled_macros_list() {
|
|
let role = Role::new("test", "---\nenabled_macros: [a, b]\n---\nPrompt");
|
|
|
|
assert_eq!(
|
|
role.enabled_macros,
|
|
Some(vec!["a".to_string(), "b".to_string()])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_enabled_macros_null_is_none() {
|
|
let role = Role::new("test", "---\nenabled_macros: null\n---\nPrompt");
|
|
|
|
assert_eq!(role.enabled_macros, None);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_parses_mcp_tools_list_and_csv_values() {
|
|
let content = "---\nmcp_tools:\n github: [get_*, list_*, search_code]\n slack: conversations_history,conversations_replies\n---\nPrompt";
|
|
|
|
let role = Role::new("test", content);
|
|
|
|
let mcp_tools = role.mcp_tools().unwrap();
|
|
assert_eq!(
|
|
mcp_tools.get("github"),
|
|
Some(&vec![
|
|
"get_*".to_string(),
|
|
"list_*".to_string(),
|
|
"search_code".to_string()
|
|
])
|
|
);
|
|
assert_eq!(
|
|
mcp_tools.get("slack"),
|
|
Some(&vec![
|
|
"conversations_history".to_string(),
|
|
"conversations_replies".to_string()
|
|
])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_mcp_tools_empty_list_server_is_some_empty() {
|
|
let role = Role::new("test", "---\nmcp_tools:\n github: []\n---\nPrompt");
|
|
|
|
assert_eq!(role.mcp_tools().unwrap().get("github"), Some(&vec![]));
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_mcp_tools_per_server_null_is_some_empty() {
|
|
let role = Role::new("test", "---\nmcp_tools:\n github:\n---\nPrompt");
|
|
|
|
assert_eq!(role.mcp_tools().unwrap().get("github"), Some(&vec![]));
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_mcp_tools_absent_is_none() {
|
|
let role = Role::new("test", "---\ntemperature: 0.5\n---\nPrompt");
|
|
|
|
assert_eq!(role.mcp_tools(), None);
|
|
}
|
|
|
|
#[test]
|
|
fn role_new_mcp_tools_null_is_none() {
|
|
let role = Role::new("test", "---\nmcp_tools: null\n---\nPrompt");
|
|
|
|
assert_eq!(role.mcp_tools(), None);
|
|
}
|
|
|
|
#[test]
|
|
fn role_export_mcp_tools_round_trips() {
|
|
let content = "---\nmcp_tools:\n github: [get_issue]\n slack: a,b\n---\nPrompt";
|
|
let role = Role::new("test", content);
|
|
|
|
let reparsed = Role::new("test", &role.export());
|
|
|
|
assert_eq!(reparsed.mcp_tools(), role.mcp_tools());
|
|
assert!(role.mcp_tools().is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn role_export_includes_enabled_macros() {
|
|
let role = Role::new("test", "---\nenabled_macros: [a]\n---\nPrompt");
|
|
|
|
let exported = role.export();
|
|
|
|
assert!(exported.contains("enabled_macros: [\"a\"]"));
|
|
}
|
|
|
|
#[test]
|
|
fn role_export_omits_enabled_macros_when_none() {
|
|
let role = Role::new("test", "Just a prompt");
|
|
|
|
assert!(!role.export().contains("enabled_macros"));
|
|
}
|
|
|
|
#[test]
|
|
fn role_builtin_shell_loads() {
|
|
let role = Role::builtin("shell").unwrap();
|
|
|
|
assert_eq!(role.name(), "shell");
|
|
assert!(!role.prompt().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn role_builtin_code_loads() {
|
|
let role = Role::builtin("code").unwrap();
|
|
|
|
assert_eq!(role.name(), "code");
|
|
assert!(!role.prompt().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn role_builtin_nonexistent_errors() {
|
|
let result = Role::builtin("nonexistent_role_xyz");
|
|
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn role_default_has_empty_fields() {
|
|
let role = Role::default();
|
|
|
|
assert_eq!(role.name(), "");
|
|
assert_eq!(role.prompt(), "");
|
|
assert_eq!(role.model_id(), None);
|
|
}
|
|
|
|
#[test]
|
|
fn role_set_model_updates_model() {
|
|
let mut role = Role::new("test", "prompt");
|
|
let model = Model::default();
|
|
|
|
role.set_model(model.clone());
|
|
|
|
assert_eq!(role.model().id(), model.id());
|
|
}
|
|
|
|
#[test]
|
|
fn role_set_temperature_works() {
|
|
let mut role = Role::new("test", "prompt");
|
|
|
|
role.set_temperature(Some(0.5));
|
|
|
|
assert_eq!(role.temperature(), Some(0.5));
|
|
}
|
|
|
|
#[test]
|
|
fn role_export_includes_metadata() {
|
|
let content = "---\ntemperature: 0.8\n---\nMy prompt";
|
|
let role = Role::new("test", content);
|
|
|
|
let exported = role.export();
|
|
|
|
assert!(exported.contains("temperature"));
|
|
assert!(exported.contains("My prompt"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_structure_prompt1() {
|
|
let prompt = r#"
|
|
System message
|
|
### INPUT:
|
|
Input 1
|
|
### OUTPUT:
|
|
Output 1
|
|
"#;
|
|
|
|
assert_eq!(
|
|
parse_structure_prompt(prompt),
|
|
("System message", vec![("Input 1", "Output 1")])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_structure_prompt2() {
|
|
let prompt = r#"
|
|
### INPUT:
|
|
Input 1
|
|
### OUTPUT:
|
|
Output 1
|
|
"#;
|
|
|
|
assert_eq!(
|
|
parse_structure_prompt(prompt),
|
|
("", vec![("Input 1", "Output 1")])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_structure_prompt3() {
|
|
let prompt = r#"
|
|
System message
|
|
### INPUT:
|
|
Input 1
|
|
"#;
|
|
|
|
assert_eq!(parse_structure_prompt(prompt), (prompt, vec![]));
|
|
}
|
|
|
|
fn fixture(entries: &[(&str, &str)]) -> Vec<(String, String)> {
|
|
entries
|
|
.iter()
|
|
.map(|(name, content)| (name.to_string(), content.to_string()))
|
|
.collect()
|
|
}
|
|
|
|
fn hooks_fixture_dir(label: &str) -> PathBuf {
|
|
let dir = temp_file(label, "");
|
|
create_dir_all(&dir).unwrap();
|
|
dir.join("hooks")
|
|
}
|
|
|
|
#[test]
|
|
fn role_hooks_install_writes_scripts_and_manifest() {
|
|
let dir = hooks_fixture_dir("role-hooks-install-");
|
|
|
|
install_and_reconcile_role_hooks(&dir, &fixture(&[("a.sh", "#!/bin/sh\n")]), false)
|
|
.unwrap();
|
|
|
|
assert_eq!(read_to_string(dir.join("a.sh")).unwrap(), "#!/bin/sh\n");
|
|
assert_eq!(
|
|
read_to_string(dir.join(builtin_manifest::BUILTIN_MANIFEST_FILE)).unwrap(),
|
|
"a.sh\n"
|
|
);
|
|
#[cfg(unix)]
|
|
{
|
|
use std::os::unix::fs::PermissionsExt;
|
|
let mode = fs::metadata(dir.join("a.sh")).unwrap().permissions().mode();
|
|
assert_eq!(mode & 0o777, 0o755, "role hook scripts install executable");
|
|
let manifest_mode = fs::metadata(dir.join(builtin_manifest::BUILTIN_MANIFEST_FILE))
|
|
.unwrap()
|
|
.permissions()
|
|
.mode();
|
|
assert_eq!(
|
|
manifest_mode & 0o111,
|
|
0,
|
|
"the manifest must never be executable"
|
|
);
|
|
}
|
|
let _ = fs::remove_dir_all(dir.parent().unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn role_hooks_install_overwrites_only_with_force() {
|
|
let dir = hooks_fixture_dir("role-hooks-force-");
|
|
let shipped = fixture(&[("a.sh", "new content\n")]);
|
|
create_dir_all(&dir).unwrap();
|
|
fs::write(dir.join("a.sh"), "SENTINEL").unwrap();
|
|
|
|
install_and_reconcile_role_hooks(&dir, &shipped, false).unwrap();
|
|
assert_eq!(read_to_string(dir.join("a.sh")).unwrap(), "SENTINEL");
|
|
|
|
install_and_reconcile_role_hooks(&dir, &shipped, true).unwrap();
|
|
assert_eq!(read_to_string(dir.join("a.sh")).unwrap(), "new content\n");
|
|
let _ = fs::remove_dir_all(dir.parent().unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn role_hooks_reconcile_removes_dropped_hook_and_keeps_user_files() {
|
|
let dir = hooks_fixture_dir("role-hooks-reconcile-");
|
|
|
|
install_and_reconcile_role_hooks(
|
|
&dir,
|
|
&fixture(&[("keep.sh", "#!/bin/sh\n"), ("drop.sh", "#!/bin/sh\n")]),
|
|
false,
|
|
)
|
|
.unwrap();
|
|
fs::write(dir.join("user.sh"), "user-owned").unwrap();
|
|
|
|
install_and_reconcile_role_hooks(&dir, &fixture(&[("keep.sh", "#!/bin/sh\n")]), false)
|
|
.unwrap();
|
|
|
|
assert!(
|
|
!dir.join("drop.sh").exists(),
|
|
"dropped builtin hook removed"
|
|
);
|
|
assert!(dir.join("keep.sh").exists());
|
|
assert!(dir.join("user.sh").exists(), "user file survives reconcile");
|
|
assert_eq!(
|
|
read_to_string(dir.join(builtin_manifest::BUILTIN_MANIFEST_FILE)).unwrap(),
|
|
"keep.sh\n"
|
|
);
|
|
let _ = fs::remove_dir_all(dir.parent().unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn role_hooks_never_claim_a_preexisting_user_file() {
|
|
let dir = hooks_fixture_dir("role-hooks-unclaimed-");
|
|
create_dir_all(&dir).unwrap();
|
|
fs::write(dir.join("a.sh"), "user-owned").unwrap();
|
|
|
|
install_and_reconcile_role_hooks(&dir, &fixture(&[("a.sh", "#!/bin/sh\n")]), false)
|
|
.unwrap();
|
|
assert!(
|
|
!dir.join(builtin_manifest::BUILTIN_MANIFEST_FILE).exists(),
|
|
"a skipped pre-existing file must not be claimed by the manifest"
|
|
);
|
|
|
|
install_and_reconcile_role_hooks(&dir, &[], false).unwrap();
|
|
assert_eq!(
|
|
read_to_string(dir.join("a.sh")).unwrap(),
|
|
"user-owned",
|
|
"the user file survives the hook being dropped from the embed"
|
|
);
|
|
let _ = fs::remove_dir_all(dir.parent().unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn role_hooks_empty_shipped_set_creates_nothing() {
|
|
let dir = hooks_fixture_dir("role-hooks-empty-");
|
|
|
|
install_and_reconcile_role_hooks(&dir, &[], false).unwrap();
|
|
|
|
assert!(!dir.exists(), "an empty shipped set must not create hooks/");
|
|
let _ = fs::remove_dir_all(dir.parent().unwrap());
|
|
}
|
|
}
|