use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use crate::error::Result;
use crate::session::SessionContext;
use crate::types::{ExtractionSource, SessionActivity, SupportTier};
pub type SessionEndCallback = Arc<dyn Fn(SessionContext) + Send + Sync>;
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct LifecycleCapabilities {
pub session_start: bool,
pub session_end: bool,
pub checkpoint: bool,
pub error_hook: bool,
pub compact: bool,
}
impl LifecycleCapabilities {
pub fn end_only() -> Self {
Self {
session_end: true,
..Default::default()
}
}
pub fn monitor_only() -> Self {
Self::default()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HookResult {
pub success: bool,
pub agent_type: String,
pub source: ExtractionSource,
pub context: Option<SessionContext>,
pub error: Option<String>,
pub timestamp: DateTime<Utc>,
}
impl HookResult {
pub fn success(agent_type: impl Into<String>, source: ExtractionSource) -> Self {
Self {
success: true,
agent_type: agent_type.into(),
source,
context: None,
error: None,
timestamp: Utc::now(),
}
}
pub fn success_with_context(
agent_type: impl Into<String>,
source: ExtractionSource,
context: SessionContext,
) -> Self {
Self {
success: true,
agent_type: agent_type.into(),
source,
context: Some(context),
error: None,
timestamp: Utc::now(),
}
}
pub fn failure(
agent_type: impl Into<String>,
source: ExtractionSource,
error: impl Into<String>,
) -> Self {
Self {
success: false,
agent_type: agent_type.into(),
source,
context: None,
error: Some(error.into()),
timestamp: Utc::now(),
}
}
}
#[async_trait]
pub trait AgentHook: Send + Sync {
fn agent_type(&self) -> &str;
async fn install_session_end_hook(&mut self, callback: SessionEndCallback) -> Result<()>;
async fn install_session_start_hook(&mut self, _callback: SessionEndCallback) -> Result<()> {
Err(crate::error::HookError::NotSupported(
"Session start hooks not supported for this agent".to_string(),
))
}
async fn install_compact_hook(&mut self, _callback: SessionEndCallback) -> Result<()> {
Err(crate::error::HookError::NotSupported(
"Compact/checkpoint hooks not supported for this agent".to_string(),
))
}
async fn detect_session_activity(&self) -> Result<SessionActivity>;
async fn extract_session_context(&self) -> Result<SessionContext>;
async fn install_checkpoint_hook(&mut self, _callback: SessionEndCallback) -> Result<()> {
Err(crate::error::HookError::NotSupported(
"Checkpoint hooks not supported for this agent".to_string(),
))
}
async fn install_error_hook(&mut self, _callback: SessionEndCallback) -> Result<()> {
Err(crate::error::HookError::NotSupported(
"Error hooks not supported for this agent".to_string(),
))
}
fn is_hook_installed(&self) -> bool {
false
}
async fn uninstall_hooks(&mut self) -> Result<()> {
Ok(())
}
fn reliability_score(&self) -> f32 {
1.0
}
fn lifecycle_capabilities(&self) -> LifecycleCapabilities {
LifecycleCapabilities::end_only()
}
fn support_tier(&self) -> SupportTier {
SupportTier::MonitorOnly
}
}
pub struct BaseHook {
pub agent_type: String,
pub installed: bool,
pub callbacks: Vec<SessionEndCallback>,
}
impl BaseHook {
pub fn new(agent_type: impl Into<String>) -> Self {
Self {
agent_type: agent_type.into(),
installed: false,
callbacks: Vec::new(),
}
}
pub fn add_callback(&mut self, callback: SessionEndCallback) {
self.callbacks.push(callback);
}
pub fn trigger_callbacks(&self, context: SessionContext) {
for callback in &self.callbacks {
callback(context.clone());
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hook_result_success() {
let result = HookResult::success("test-agent", ExtractionSource::Manual);
assert!(result.success);
assert!(result.error.is_none());
}
#[test]
fn test_hook_result_failure() {
let result = HookResult::failure(
"test-agent",
ExtractionSource::Manual,
"Something went wrong",
);
assert!(!result.success);
assert!(result.error.is_some());
assert_eq!(result.error.unwrap(), "Something went wrong");
}
#[test]
fn test_hook_result_with_context() {
let ctx = SessionContext::new("test");
let result = HookResult::success_with_context(
"test-agent",
ExtractionSource::NativeHook("skill".to_string()),
ctx,
);
assert!(result.success);
assert!(result.context.is_some());
}
#[test]
fn test_base_hook() {
let mut hook = BaseHook::new("test");
assert_eq!(hook.agent_type, "test");
assert!(!hook.installed);
let called = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let called_clone = called.clone();
hook.add_callback(Arc::new(move |_ctx| {
called_clone.store(true, std::sync::atomic::Ordering::SeqCst);
}));
hook.trigger_callbacks(SessionContext::new("test"));
assert!(called.load(std::sync::atomic::Ordering::SeqCst));
}
#[test]
fn test_lifecycle_capabilities_default() {
let caps = LifecycleCapabilities::default();
assert!(!caps.session_start);
assert!(!caps.session_end);
assert!(!caps.checkpoint);
assert!(!caps.error_hook);
assert!(!caps.compact);
}
#[test]
fn test_lifecycle_capabilities_end_only() {
let caps = LifecycleCapabilities::end_only();
assert!(!caps.session_start);
assert!(caps.session_end);
assert!(!caps.checkpoint);
assert!(!caps.error_hook);
assert!(!caps.compact);
}
#[test]
fn test_lifecycle_capabilities_monitor_only() {
let caps = LifecycleCapabilities::monitor_only();
assert!(!caps.session_end);
assert!(!caps.session_start);
}
}