mcp_usage_tower/
classify.rs1use std::borrow::Cow;
4use std::fmt;
5
6use mcp_usage_core::{Method, name, validate_header};
7
8#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct ProtocolHeaders<'a> {
11 pub method: Method,
13 pub name: Option<Cow<'a, str>>,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq)]
19pub enum ClassificationError {
20 UntrustedVersion(String),
22 MissingMethod,
24 MissingName,
26 InvalidName(String),
28}
29
30impl fmt::Display for ClassificationError {
31 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32 match self {
33 Self::UntrustedVersion(reason) => write!(f, "untrusted MCP headers: {reason}"),
34 Self::MissingMethod => f.write_str("missing Mcp-Method header"),
35 Self::MissingName => f.write_str("missing Mcp-Name header for name-bearing method"),
36 Self::InvalidName(reason) => write!(f, "invalid Mcp-Name header: {reason}"),
37 }
38 }
39}
40
41impl std::error::Error for ClassificationError {}
42
43pub fn classify_protocol_headers<'a>(
52 version_header: Option<&str>,
53 method_header: Option<&str>,
54 name_header: Option<&'a str>,
55) -> Result<ProtocolHeaders<'a>, ClassificationError> {
56 if let Err(reason) = validate_header(version_header) {
57 return Err(ClassificationError::UntrustedVersion(reason.to_string()));
58 }
59 let method = Method::parse(method_header.ok_or(ClassificationError::MissingMethod)?);
60 if method.carries_name() && name_header.is_none() {
61 return Err(ClassificationError::MissingName);
62 }
63 let name = name_header
64 .map(name::decode)
65 .transpose()
66 .map_err(|error| ClassificationError::InvalidName(error.to_string()))?;
67 Ok(ProtocolHeaders { method, name })
68}
69
70#[cfg(test)]
71mod tests {
72 use super::*;
73
74 #[test]
75 fn common_path_borrows_the_name() {
76 let classified =
77 classify_protocol_headers(Some("2026-07-28"), Some("tools/call"), Some("weather"))
78 .unwrap();
79 assert_eq!(classified.method, Method::ToolsCall);
80 assert!(matches!(classified.name, Some(Cow::Borrowed("weather"))));
81 }
82
83 #[test]
84 fn legacy_requests_are_rejected_before_headers_are_trusted() {
85 let error =
86 classify_protocol_headers(Some("2025-11-25"), Some("tools/call"), Some("cheap_tool"))
87 .unwrap_err();
88 assert!(matches!(error, ClassificationError::UntrustedVersion(_)));
89 }
90
91 #[test]
92 fn name_is_required_only_where_the_spec_requires_it() {
93 assert!(matches!(
94 classify_protocol_headers(Some("2026-07-28"), Some("tools/call"), None),
95 Err(ClassificationError::MissingName)
96 ));
97 assert!(classify_protocol_headers(Some("2026-07-28"), Some("tools/list"), None).is_ok());
98 }
99
100 #[test]
101 fn header_hot_path_has_no_body_dependency() {
102 let invalid_body = b"this is not JSON";
105 let classified =
106 classify_protocol_headers(Some("2026-07-28"), Some("tools/call"), Some("weather"))
107 .unwrap();
108 assert_eq!(classified.method, Method::ToolsCall);
109 assert_eq!(invalid_body, b"this is not JSON");
110 }
111}