use std::collections::HashMap;
use std::sync::{OnceLock, RwLock};
use crate::cmf::constants::{
ENTITY_LLM, ENTITY_PROMPT, ENTITY_RESOURCE, ENTITY_TOOL, HOOK_CMF_LLM_INPUT,
HOOK_CMF_LLM_OUTPUT, HOOK_CMF_PROMPT_POST_INVOKE, HOOK_CMF_PROMPT_PRE_INVOKE,
HOOK_CMF_RESOURCE_POST_FETCH, HOOK_CMF_RESOURCE_PRE_FETCH, HOOK_CMF_TOOL_POST_INVOKE,
HOOK_CMF_TOOL_PRE_INVOKE,
};
use crate::delegation::HOOK_TOKEN_DELEGATE;
use crate::identity::HOOK_IDENTITY_RESOLVE;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum HookPhase {
Pre,
Post,
Unphased,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct HookMetadata {
pub entity_type: Option<&'static str>,
pub phase: HookPhase,
}
impl HookMetadata {
pub const fn unknown() -> Self {
Self {
entity_type: None,
phase: HookPhase::Unphased,
}
}
pub fn matches(&self, request_entity_type: Option<&str>, requested_phase: HookPhase) -> bool {
let entity_ok = match (self.entity_type, request_entity_type) {
(Some(hook_et), Some(req_et)) => hook_et == req_et,
(Some(_), None) => true, (None, _) => true, };
if !entity_ok {
return false;
}
match (self.phase, requested_phase) {
(HookPhase::Unphased, _) | (_, HookPhase::Unphased) => true,
(a, b) => a == b,
}
}
}
const BUILTIN_METADATA: &[(&str, HookMetadata)] = &[
(
HOOK_CMF_TOOL_PRE_INVOKE,
HookMetadata {
entity_type: Some(ENTITY_TOOL),
phase: HookPhase::Pre,
},
),
(
HOOK_CMF_TOOL_POST_INVOKE,
HookMetadata {
entity_type: Some(ENTITY_TOOL),
phase: HookPhase::Post,
},
),
(
HOOK_CMF_LLM_INPUT,
HookMetadata {
entity_type: Some(ENTITY_LLM),
phase: HookPhase::Pre,
},
),
(
HOOK_CMF_LLM_OUTPUT,
HookMetadata {
entity_type: Some(ENTITY_LLM),
phase: HookPhase::Post,
},
),
(
HOOK_CMF_PROMPT_PRE_INVOKE,
HookMetadata {
entity_type: Some(ENTITY_PROMPT),
phase: HookPhase::Pre,
},
),
(
HOOK_CMF_PROMPT_POST_INVOKE,
HookMetadata {
entity_type: Some(ENTITY_PROMPT),
phase: HookPhase::Post,
},
),
(
HOOK_CMF_RESOURCE_PRE_FETCH,
HookMetadata {
entity_type: Some(ENTITY_RESOURCE),
phase: HookPhase::Pre,
},
),
(
HOOK_CMF_RESOURCE_POST_FETCH,
HookMetadata {
entity_type: Some(ENTITY_RESOURCE),
phase: HookPhase::Post,
},
),
(
HOOK_IDENTITY_RESOLVE,
HookMetadata {
entity_type: None,
phase: HookPhase::Unphased,
},
),
(
HOOK_TOKEN_DELEGATE,
HookMetadata {
entity_type: None,
phase: HookPhase::Unphased,
},
),
];
fn registry() -> &'static RwLock<HashMap<String, HookMetadata>> {
static REGISTRY: OnceLock<RwLock<HashMap<String, HookMetadata>>> = OnceLock::new();
REGISTRY.get_or_init(|| {
let mut map: HashMap<String, HookMetadata> = HashMap::new();
for (name, meta) in BUILTIN_METADATA {
map.insert((*name).to_string(), *meta);
}
RwLock::new(map)
})
}
pub fn lookup(hook_name: &str) -> HookMetadata {
let r = registry().read().unwrap_or_else(|p| p.into_inner());
r.get(hook_name).copied().unwrap_or(HookMetadata::unknown())
}
pub fn register_hook_metadata(hook_name: impl Into<String>, meta: HookMetadata) {
let mut w = registry().write().unwrap_or_else(|p| p.into_inner());
w.insert(hook_name.into(), meta);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cmf_tool_pre_invoke_is_pre_phase_for_tool_entity() {
let meta = lookup(HOOK_CMF_TOOL_PRE_INVOKE);
assert_eq!(meta.entity_type, Some(ENTITY_TOOL));
assert_eq!(meta.phase, HookPhase::Pre);
}
#[test]
fn cmf_llm_output_is_post_phase_for_llm_entity() {
let meta = lookup(HOOK_CMF_LLM_OUTPUT);
assert_eq!(meta.entity_type, Some(ENTITY_LLM));
assert_eq!(meta.phase, HookPhase::Post);
}
#[test]
fn identity_resolve_is_unphased_no_entity() {
let meta = lookup(HOOK_IDENTITY_RESOLVE);
assert_eq!(meta.entity_type, None);
assert_eq!(meta.phase, HookPhase::Unphased);
}
#[test]
fn token_delegate_is_unphased_no_entity() {
let meta = lookup(HOOK_TOKEN_DELEGATE);
assert_eq!(meta.entity_type, None);
assert_eq!(meta.phase, HookPhase::Unphased);
}
#[test]
fn unknown_hook_returns_universal_default() {
let meta = lookup("custom.unrecognized_hook");
assert_eq!(meta.entity_type, None);
assert_eq!(meta.phase, HookPhase::Unphased);
}
#[test]
fn matches_filters_by_entity_type_when_set() {
let tool_pre = HookMetadata {
entity_type: Some(ENTITY_TOOL),
phase: HookPhase::Pre,
};
assert!(tool_pre.matches(Some(ENTITY_TOOL), HookPhase::Pre));
assert!(!tool_pre.matches(Some(ENTITY_LLM), HookPhase::Pre));
}
#[test]
fn matches_allows_any_entity_when_hook_entity_is_none() {
let universal = HookMetadata {
entity_type: None,
phase: HookPhase::Pre,
};
assert!(universal.matches(Some(ENTITY_TOOL), HookPhase::Pre));
assert!(universal.matches(Some(ENTITY_LLM), HookPhase::Pre));
assert!(universal.matches(None, HookPhase::Pre));
}
#[test]
fn matches_phase_exactly_unless_unphased() {
let tool_pre = HookMetadata {
entity_type: Some(ENTITY_TOOL),
phase: HookPhase::Pre,
};
assert!(tool_pre.matches(Some(ENTITY_TOOL), HookPhase::Pre));
assert!(!tool_pre.matches(Some(ENTITY_TOOL), HookPhase::Post));
}
#[test]
fn matches_unphased_is_wildcard_in_either_direction() {
let unphased = HookMetadata {
entity_type: None,
phase: HookPhase::Unphased,
};
assert!(unphased.matches(Some(ENTITY_TOOL), HookPhase::Pre));
assert!(unphased.matches(Some(ENTITY_LLM), HookPhase::Post));
let tool_pre = HookMetadata {
entity_type: Some(ENTITY_TOOL),
phase: HookPhase::Pre,
};
assert!(tool_pre.matches(Some(ENTITY_TOOL), HookPhase::Unphased));
}
#[test]
fn matches_request_without_entity_type_doesnt_filter_on_it() {
let tool_pre = HookMetadata {
entity_type: Some(ENTITY_TOOL),
phase: HookPhase::Pre,
};
assert!(tool_pre.matches(None, HookPhase::Pre));
}
#[test]
fn register_hook_metadata_overrides_default() {
let name = "test_custom.overridden_meta";
register_hook_metadata(
name,
HookMetadata {
entity_type: Some("custom"),
phase: HookPhase::Pre,
},
);
let meta = lookup(name);
assert_eq!(meta.entity_type, Some("custom"));
assert_eq!(meta.phase, HookPhase::Pre);
}
}