use std::collections::HashMap;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use crate::client::{ModelConfig, ModelType};
use crate::common::{ContentSource, Index, Named, SourceType, ToolRestricted};
use crate::hooks::HookRule;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SubagentIndex {
pub name: String,
pub description: String,
#[serde(default, alias = "tools")]
pub allowed_tools: Vec<String>,
pub source: ContentSource,
#[serde(default)]
pub source_type: SourceType,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model_type: Option<ModelType>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub skills: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub disallowed_tools: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub permission_mode: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hooks: Option<HashMap<String, Vec<HookRule>>>,
}
impl SubagentIndex {
pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
Self {
name: name.into(),
description: description.into(),
allowed_tools: Vec::new(),
source: ContentSource::default(),
source_type: SourceType::default(),
model: None,
model_type: None,
skills: Vec::new(),
disallowed_tools: Vec::new(),
permission_mode: None,
hooks: None,
}
}
pub fn allowed_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.allowed_tools = tools.into_iter().map(Into::into).collect();
self
}
pub fn tools(self, t: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.allowed_tools(t)
}
pub fn source(mut self, source: ContentSource) -> Self {
self.source = source;
self
}
pub fn source_type(mut self, source_type: SourceType) -> Self {
self.source_type = source_type;
self
}
pub fn model(mut self, model: impl Into<String>) -> Self {
self.model = Some(model.into());
self
}
pub fn model_type(mut self, model_type: ModelType) -> Self {
self.model_type = Some(model_type);
self
}
pub fn skills(mut self, skills: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.skills = skills.into_iter().map(Into::into).collect();
self
}
pub fn resolve_model<'a>(&'a self, config: &'a ModelConfig) -> &'a str {
if let Some(ref model) = self.model {
return config.resolve_alias(model);
}
config.get(self.model_type.unwrap_or_default())
}
pub async fn load_prompt(&self) -> crate::Result<String> {
self.source.load().await
}
}
impl Named for SubagentIndex {
fn name(&self) -> &str {
&self.name
}
}
impl ToolRestricted for SubagentIndex {
fn allowed_tools(&self) -> &[String] {
&self.allowed_tools
}
}
#[async_trait]
impl Index for SubagentIndex {
fn source(&self) -> &ContentSource {
&self.source
}
fn source_type(&self) -> SourceType {
self.source_type
}
fn to_summary_line(&self) -> String {
let tools_str = if self.allowed_tools.is_empty() {
"*".to_string()
} else {
self.allowed_tools.join(", ")
};
format!(
"- {}: {} (Tools: {})",
self.name, self.description, tools_str
)
}
fn description(&self) -> &str {
&self.description
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_subagent_index_creation() {
let subagent = SubagentIndex::new("reviewer", "Code reviewer")
.source(ContentSource::in_memory("Review the code"))
.source_type(SourceType::Project)
.tools(["Read", "Grep", "Glob"])
.model("haiku");
assert_eq!(subagent.name, "reviewer");
assert!(subagent.has_tool_restrictions());
assert!(subagent.is_tool_allowed("Read"));
assert!(!subagent.is_tool_allowed("Bash"));
}
#[test]
fn test_summary_line() {
let subagent = SubagentIndex::new("Explore", "Fast codebase exploration")
.tools(["Read", "Grep", "Glob", "Bash"]);
let summary = subagent.to_summary_line();
assert!(summary.contains("Explore"));
assert!(summary.contains("Fast codebase exploration"));
assert!(summary.contains("Read, Grep, Glob, Bash"));
}
#[test]
fn test_summary_line_no_tools() {
let subagent = SubagentIndex::new("general-purpose", "General purpose agent");
let summary = subagent.to_summary_line();
assert!(summary.contains("(Tools: *)"));
}
#[tokio::test]
async fn test_load_prompt() {
let subagent = SubagentIndex::new("test", "Test agent")
.source(ContentSource::in_memory("You are a test agent."));
let prompt = subagent.load_prompt().await.unwrap();
assert_eq!(prompt, "You are a test agent.");
}
#[test]
fn test_resolve_model_with_alias() {
let config = ModelConfig::default();
let subagent = SubagentIndex::new("fast", "Fast agent")
.source(ContentSource::in_memory("Be quick"))
.model("haiku");
assert!(subagent.resolve_model(&config).contains("haiku"));
let subagent = SubagentIndex::new("smart", "Smart agent")
.source(ContentSource::in_memory("Think deep"))
.model("opus");
assert!(subagent.resolve_model(&config).contains("opus"));
}
#[test]
fn test_resolve_model_with_type() {
let config = ModelConfig::default();
let subagent = SubagentIndex::new("typed", "Typed agent")
.source(ContentSource::in_memory("Use type"))
.model_type(ModelType::Small);
assert!(subagent.resolve_model(&config).contains("haiku"));
}
}