use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
#[repr(u8)]
pub enum ContentTrustLevel {
Trusted = 0,
LocalUntrusted = 1,
ExternalUntrusted = 2,
}
impl ContentTrustLevel {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::Trusted => "trusted",
Self::LocalUntrusted => "local_untrusted",
Self::ExternalUntrusted => "external_untrusted",
}
}
#[must_use]
pub fn from_str_opt(s: &str) -> Option<Self> {
match s {
"trusted" => Some(Self::Trusted),
"local_untrusted" => Some(Self::LocalUntrusted),
"external_untrusted" => Some(Self::ExternalUntrusted),
_ => None,
}
}
#[must_use]
pub fn from_ordinal(ordinal: u8) -> Self {
match ordinal {
0 => Self::Trusted,
1 => Self::LocalUntrusted,
_ => Self::ExternalUntrusted,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum ContentSourceKind {
ToolResult,
WebScrape,
McpResponse,
A2aMessage,
MemoryRetrieval,
InstructionFile,
ChannelMessage,
}
impl ContentSourceKind {
#[must_use]
pub fn default_trust_level(self) -> ContentTrustLevel {
match self {
Self::ToolResult | Self::InstructionFile => ContentTrustLevel::LocalUntrusted,
Self::WebScrape
| Self::McpResponse
| Self::A2aMessage
| Self::MemoryRetrieval
| Self::ChannelMessage => ContentTrustLevel::ExternalUntrusted,
}
}
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::ToolResult => "tool_result",
Self::WebScrape => "web_scrape",
Self::McpResponse => "mcp_response",
Self::A2aMessage => "a2a_message",
Self::MemoryRetrieval => "memory_retrieval",
Self::InstructionFile => "instruction_file",
Self::ChannelMessage => "channel_message",
}
}
#[must_use]
pub fn from_str_opt(s: &str) -> Option<Self> {
match s {
"tool_result" => Some(Self::ToolResult),
"web_scrape" => Some(Self::WebScrape),
"mcp_response" => Some(Self::McpResponse),
"a2a_message" => Some(Self::A2aMessage),
"memory_retrieval" => Some(Self::MemoryRetrieval),
"instruction_file" => Some(Self::InstructionFile),
"channel_message" => Some(Self::ChannelMessage),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum MemorySourceHint {
ConversationHistory,
LlmSummary,
ExternalContent,
}
#[derive(Debug, Clone)]
pub struct ContentSource {
pub kind: ContentSourceKind,
pub trust_level: ContentTrustLevel,
pub identifier: Option<String>,
pub memory_hint: Option<MemorySourceHint>,
}
impl ContentSource {
#[must_use]
pub fn new(kind: ContentSourceKind) -> Self {
Self {
trust_level: kind.default_trust_level(),
kind,
identifier: None,
memory_hint: None,
}
}
#[must_use]
pub fn with_identifier(mut self, id: impl Into<String>) -> Self {
self.identifier = Some(id.into());
self
}
#[must_use]
pub fn with_trust_level(mut self, level: ContentTrustLevel) -> Self {
self.trust_level = level;
self
}
#[must_use]
pub fn with_memory_hint(mut self, hint: MemorySourceHint) -> Self {
self.memory_hint = Some(hint);
self
}
}
#[derive(Debug, Clone)]
pub struct InjectionFlag {
pub pattern_name: &'static str,
pub byte_offset: usize,
pub matched_text: String,
}
#[cfg(feature = "classifiers")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum InjectionVerdict {
Clean,
Suspicious,
Blocked,
}
#[cfg(feature = "classifiers")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum InstructionClass {
NoInstruction,
AlignedInstruction,
MisalignedInstruction,
Unknown,
}
#[cfg(feature = "classifiers")]
impl InstructionClass {
pub(crate) fn from_label(label: &str) -> Self {
match label.to_lowercase().as_str() {
"no_instruction" | "no-instruction" | "none" => Self::NoInstruction,
"aligned_instruction" | "aligned-instruction" | "aligned" => Self::AlignedInstruction,
"misaligned_instruction" | "misaligned-instruction" | "misaligned" => {
Self::MisalignedInstruction
}
_ => Self::Unknown,
}
}
}
#[derive(Debug, Clone)]
pub struct SanitizedContent {
pub body: String,
pub source: ContentSource,
pub injection_flags: Vec<InjectionFlag>,
pub was_truncated: bool,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn from_str_opt_known_variants_roundtrip() {
let variants = [
(ContentSourceKind::ToolResult, "tool_result"),
(ContentSourceKind::WebScrape, "web_scrape"),
(ContentSourceKind::McpResponse, "mcp_response"),
(ContentSourceKind::A2aMessage, "a2a_message"),
(ContentSourceKind::MemoryRetrieval, "memory_retrieval"),
(ContentSourceKind::InstructionFile, "instruction_file"),
(ContentSourceKind::ChannelMessage, "channel_message"),
];
for (kind, s) in &variants {
assert_eq!(ContentSourceKind::from_str_opt(s), Some(*kind));
assert_eq!(kind.as_str(), *s);
}
}
#[test]
fn from_str_opt_unknown_returns_none() {
assert_eq!(ContentSourceKind::from_str_opt("unknown"), None);
assert_eq!(ContentSourceKind::from_str_opt(""), None);
}
#[test]
fn from_str_opt_case_sensitive() {
assert_eq!(ContentSourceKind::from_str_opt("WebScrape"), None);
assert_eq!(ContentSourceKind::from_str_opt("TOOL_RESULT"), None);
}
#[test]
fn content_source_new_has_default_trust_and_no_identifier() {
let source = ContentSource::new(ContentSourceKind::WebScrape);
assert_eq!(source.trust_level, ContentTrustLevel::ExternalUntrusted);
assert!(source.identifier.is_none());
assert!(source.memory_hint.is_none());
}
#[test]
fn content_source_with_identifier() {
let source = ContentSource::new(ContentSourceKind::ToolResult).with_identifier("shell");
assert_eq!(source.identifier.as_deref(), Some("shell"));
}
#[test]
fn content_source_with_trust_level_override() {
let source = ContentSource::new(ContentSourceKind::McpResponse)
.with_trust_level(ContentTrustLevel::LocalUntrusted);
assert_eq!(source.trust_level, ContentTrustLevel::LocalUntrusted);
}
#[test]
fn content_source_with_memory_hint() {
let source = ContentSource::new(ContentSourceKind::MemoryRetrieval)
.with_memory_hint(MemorySourceHint::ConversationHistory);
assert_eq!(
source.memory_hint,
Some(MemorySourceHint::ConversationHistory)
);
}
#[test]
fn content_trust_level_equality() {
assert_eq!(ContentTrustLevel::Trusted, ContentTrustLevel::Trusted);
assert_ne!(
ContentTrustLevel::Trusted,
ContentTrustLevel::LocalUntrusted
);
assert_ne!(
ContentTrustLevel::LocalUntrusted,
ContentTrustLevel::ExternalUntrusted
);
}
#[test]
fn trust_level_from_str_opt_known_variants_roundtrip() {
let variants = [
(ContentTrustLevel::Trusted, "trusted"),
(ContentTrustLevel::LocalUntrusted, "local_untrusted"),
(ContentTrustLevel::ExternalUntrusted, "external_untrusted"),
];
for (level, s) in &variants {
assert_eq!(ContentTrustLevel::from_str_opt(s), Some(*level));
assert_eq!(level.as_str(), *s);
}
}
#[test]
fn trust_level_from_str_opt_unknown_returns_none() {
assert_eq!(ContentTrustLevel::from_str_opt("unknown"), None);
assert_eq!(ContentTrustLevel::from_str_opt(""), None);
}
#[test]
fn trust_level_ord_matches_severity() {
assert!(ContentTrustLevel::Trusted < ContentTrustLevel::LocalUntrusted);
assert!(ContentTrustLevel::LocalUntrusted < ContentTrustLevel::ExternalUntrusted);
assert_eq!(
ContentTrustLevel::Trusted.max(ContentTrustLevel::ExternalUntrusted),
ContentTrustLevel::ExternalUntrusted
);
}
#[test]
fn default_trust_level_local_kinds() {
assert_eq!(
ContentSourceKind::ToolResult.default_trust_level(),
ContentTrustLevel::LocalUntrusted
);
assert_eq!(
ContentSourceKind::InstructionFile.default_trust_level(),
ContentTrustLevel::LocalUntrusted
);
}
#[test]
fn default_trust_level_external_kinds() {
for kind in [
ContentSourceKind::WebScrape,
ContentSourceKind::McpResponse,
ContentSourceKind::A2aMessage,
ContentSourceKind::MemoryRetrieval,
ContentSourceKind::ChannelMessage,
] {
assert_eq!(
kind.default_trust_level(),
ContentTrustLevel::ExternalUntrusted
);
}
}
}