use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::BTreeMap;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct OffdeviceCallDescriptor {
pub kind: String,
pub provider: String,
pub model: String,
pub subsystem: String,
#[serde(default)]
pub metadata: BTreeMap<String, Value>,
}
impl OffdeviceCallDescriptor {
pub fn new(
kind: impl Into<String>,
provider: impl Into<String>,
model: impl Into<String>,
subsystem: impl Into<String>,
) -> Self {
Self {
kind: kind.into(),
provider: provider.into(),
model: model.into(),
subsystem: subsystem.into(),
metadata: BTreeMap::new(),
}
}
pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
self.metadata.insert(key.into(), value.into());
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum OffdeviceDecision {
Allow,
Veto {
reason: String,
},
}
#[async_trait]
pub trait OffdeviceCallHook: Send + Sync {
async fn before_call(&self, _descriptor: &OffdeviceCallDescriptor) -> OffdeviceDecision {
OffdeviceDecision::Allow
}
}
pub struct ClosureOffdeviceHook<F>(pub F)
where
F: Fn(&OffdeviceCallDescriptor) -> OffdeviceDecision + Send + Sync + 'static;
#[async_trait]
impl<F> OffdeviceCallHook for ClosureOffdeviceHook<F>
where
F: Fn(&OffdeviceCallDescriptor) -> OffdeviceDecision + Send + Sync + 'static,
{
async fn before_call(&self, descriptor: &OffdeviceCallDescriptor) -> OffdeviceDecision {
(self.0)(descriptor)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn closure_hook_vetoes_anthropic_calls() {
let hook = ClosureOffdeviceHook(|d: &OffdeviceCallDescriptor| {
if d.provider == "anthropic" {
OffdeviceDecision::Veto {
reason: "no anthropic calls".into(),
}
} else {
OffdeviceDecision::Allow
}
});
let d = OffdeviceCallDescriptor::new(
"summarizer",
"anthropic",
"claude-haiku-4-5",
"memory_summarizer",
);
match hook.before_call(&d).await {
OffdeviceDecision::Veto { reason } => assert_eq!(reason, "no anthropic calls"),
other => panic!("expected Veto, got {other:?}"),
}
let d2 =
OffdeviceCallDescriptor::new("summarizer", "ollama", "llama3.2", "memory_summarizer");
assert_eq!(hook.before_call(&d2).await, OffdeviceDecision::Allow);
}
#[tokio::test]
async fn default_hook_allows_everything() {
struct Default;
#[async_trait]
impl OffdeviceCallHook for Default {}
let d = OffdeviceCallDescriptor::new("summarizer", "anthropic", "haiku", "memory");
assert_eq!(Default.before_call(&d).await, OffdeviceDecision::Allow);
}
#[tokio::test]
async fn descriptor_metadata_builder() {
let d = OffdeviceCallDescriptor::new("summarizer", "anthropic", "haiku", "memory")
.with_metadata("session_id", "s1")
.with_metadata("chunk_count", 5);
assert_eq!(d.metadata.len(), 2);
assert_eq!(
d.metadata.get("session_id").and_then(|v| v.as_str()),
Some("s1")
);
assert_eq!(
d.metadata.get("chunk_count").and_then(|v| v.as_i64()),
Some(5)
);
}
}