use std::fmt;
use std::sync::Arc;
use async_trait::async_trait;
use uuid::Uuid;
use crate::context::{InvestigationContext, NextAction};
use crate::registry::{KernelError, SkillRegistry, ToolRegistry};
use crate::skill::{Skill, SkillOutcome};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct AgentId(pub Uuid);
impl AgentId {
pub fn new() -> Self {
Self(Uuid::new_v4())
}
}
impl Default for AgentId {
fn default() -> Self {
Self::new()
}
}
impl fmt::Display for AgentId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Debug, Clone)]
pub struct AgentStepResult {
pub skills_run: Vec<String>,
pub skills_skipped: Vec<String>,
pub confidence: f32,
pub concluded: bool,
}
#[async_trait]
pub trait AgentLifecycleHook: Send + Sync {
async fn before_step(
&self,
_agent: &GenericAgent,
_ctx: &InvestigationContext,
) -> Result<(), KernelError> {
Ok(())
}
async fn before_skill(
&self,
_agent: &GenericAgent,
_skill_id: &str,
_applies: bool,
_confidence: f32,
) -> Result<(), KernelError> {
Ok(())
}
async fn after_skill(
&self,
_agent: &GenericAgent,
_skill_id: &str,
_outcome: &SkillOutcome,
_confidence: f32,
) -> Result<(), KernelError> {
Ok(())
}
async fn after_step(
&self,
_agent: &GenericAgent,
_result: &AgentStepResult,
) -> Result<(), KernelError> {
Ok(())
}
async fn on_step_error(
&self,
_agent: &GenericAgent,
_error: &KernelError,
) -> Result<(), KernelError> {
Ok(())
}
}
#[async_trait]
pub trait Agent: Send + Sync {
fn id(&self) -> AgentId;
fn name(&self) -> &str;
async fn step(&self, ctx: &mut InvestigationContext) -> Result<AgentStepResult, KernelError>;
}
pub struct GenericAgent {
id: AgentId,
name: String,
skills: Vec<Arc<dyn Skill>>,
tools: ToolRegistry,
lifecycle_hooks: Vec<Arc<dyn AgentLifecycleHook>>,
pub short_circuit_on_conclude: bool,
}
impl GenericAgent {
pub fn builder(name: impl Into<String>) -> GenericAgentBuilder {
GenericAgentBuilder {
name: name.into(),
skill_ids: Vec::new(),
allowed_tools: None,
lifecycle_hooks: Vec::new(),
short_circuit_on_conclude: true,
}
}
pub fn skills(&self) -> &[Arc<dyn Skill>] {
&self.skills
}
pub fn tools(&self) -> &ToolRegistry {
&self.tools
}
}
#[async_trait]
impl Agent for GenericAgent {
fn id(&self) -> AgentId {
self.id
}
fn name(&self) -> &str {
&self.name
}
async fn step(&self, ctx: &mut InvestigationContext) -> Result<AgentStepResult, KernelError> {
let mut skills_run = Vec::new();
let mut skills_skipped = Vec::new();
let mut concluded = false;
self.notify_before_step(ctx).await;
for skill in &self.skills {
if concluded && self.short_circuit_on_conclude {
let skill_id = skill.id();
self.notify_before_skill(skill_id, false, ctx.confidence)
.await;
skills_skipped.push(skill.id().to_string());
continue;
}
if !skill.applies(ctx) {
let skill_id = skill.id();
self.notify_before_skill(skill_id, false, ctx.confidence)
.await;
skills_skipped.push(skill.id().to_string());
continue;
}
let skill_id = skill.id();
self.notify_before_skill(skill_id, true, ctx.confidence)
.await;
let outcome = match skill.execute(ctx, &self.tools).await {
Ok(outcome) => outcome,
Err(error) => {
self.notify_step_error(&error).await;
return Err(error);
}
};
ctx.confidence = (ctx.confidence + outcome.confidence_delta).clamp(0.0, 1.0);
ctx.pending_actions = outcome.next_actions.clone();
self.notify_after_skill(skill_id, &outcome, ctx.confidence)
.await;
if outcome
.next_actions
.iter()
.any(|a| matches!(a, NextAction::Conclude | NextAction::Discard))
{
concluded = true;
}
skills_run.push(skill_id.to_string());
}
let result = AgentStepResult {
skills_run,
skills_skipped,
confidence: ctx.confidence,
concluded,
};
self.notify_after_step(&result).await;
Ok(result)
}
}
impl GenericAgent {
async fn notify_before_step(&self, ctx: &InvestigationContext) {
for hook in &self.lifecycle_hooks {
if let Err(err) = hook.before_step(self, ctx).await {
tracing::warn!(error = %err, "lifecycle hook before_step failed; continuing");
}
}
}
async fn notify_before_skill(&self, skill_id: &str, applies: bool, confidence: f32) {
for hook in &self.lifecycle_hooks {
if let Err(err) = hook.before_skill(self, skill_id, applies, confidence).await {
tracing::warn!(
error = %err,
skill_id,
"lifecycle hook before_skill failed; continuing"
);
}
}
}
async fn notify_after_skill(&self, skill_id: &str, outcome: &SkillOutcome, confidence: f32) {
for hook in &self.lifecycle_hooks {
if let Err(err) = hook.after_skill(self, skill_id, outcome, confidence).await {
tracing::warn!(
error = %err,
skill_id,
"lifecycle hook after_skill failed; continuing"
);
}
}
}
async fn notify_after_step(&self, result: &AgentStepResult) {
for hook in &self.lifecycle_hooks {
if let Err(err) = hook.after_step(self, result).await {
tracing::warn!(error = %err, "lifecycle hook after_step failed; continuing");
}
}
}
async fn notify_step_error(&self, error: &KernelError) {
for hook in &self.lifecycle_hooks {
if let Err(err) = hook.on_step_error(self, error).await {
tracing::warn!(error = %err, "lifecycle hook on_step_error failed; continuing");
}
}
}
}
pub struct GenericAgentBuilder {
name: String,
skill_ids: Vec<String>,
allowed_tools: Option<Vec<String>>,
lifecycle_hooks: Vec<Arc<dyn AgentLifecycleHook>>,
short_circuit_on_conclude: bool,
}
impl GenericAgentBuilder {
pub fn with_skills<I, S>(mut self, ids: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.skill_ids.extend(ids.into_iter().map(Into::into));
self
}
pub fn with_tools<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.allowed_tools = Some(names.into_iter().map(Into::into).collect());
self
}
pub fn short_circuit_on_conclude(mut self, v: bool) -> Self {
self.short_circuit_on_conclude = v;
self
}
pub fn with_lifecycle_hook(mut self, hook: Arc<dyn AgentLifecycleHook>) -> Self {
self.lifecycle_hooks.push(hook);
self
}
pub fn build(
self,
skills: &SkillRegistry,
tools: &ToolRegistry,
) -> Result<GenericAgent, KernelError> {
let resolved = skills.resolve_chain(self.skill_ids.iter())?;
let scoped_tools = match self.allowed_tools {
Some(list) => tools.scoped(list),
None => tools.clone(),
};
Ok(GenericAgent {
id: AgentId::new(),
name: self.name,
skills: resolved,
tools: scoped_tools,
lifecycle_hooks: self.lifecycle_hooks,
short_circuit_on_conclude: self.short_circuit_on_conclude,
})
}
}