Alex Hancock
109 lines
3.8 KiB
Rust
109 lines
3.8 KiB
Rust
use base64::Engine;
|
|
pub use rmcp::model::ErrorData;
|
|
use rmcp::model::ResourceContents;
|
|
|
|
pub type ToolResult<T> = Result<T, ErrorData>;
|
|
|
|
pub fn extract_text_from_resource(resource: &ResourceContents) -> String {
|
|
match resource {
|
|
ResourceContents::TextResourceContents { text, .. } => text.clone(),
|
|
ResourceContents::BlobResourceContents {
|
|
blob, mime_type, ..
|
|
} => match base64::engine::general_purpose::STANDARD.decode(blob) {
|
|
Ok(bytes) => {
|
|
let byte_len = bytes.len();
|
|
match String::from_utf8(bytes) {
|
|
Ok(text) => text,
|
|
Err(_) => {
|
|
let mime = mime_type
|
|
.as_ref()
|
|
.map(|m| m.as_str())
|
|
.unwrap_or("application/octet-stream");
|
|
format!("[Binary content ({}) - {} bytes]", mime, byte_len)
|
|
}
|
|
}
|
|
}
|
|
Err(_) => blob.clone(),
|
|
},
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use test_case::test_case;
|
|
|
|
#[test_case("Hello, World!", "Hello, World!" ; "simple text")]
|
|
#[test_case("Hello from GitHub!", "Hello from GitHub!" ; "github content")]
|
|
#[test_case("", "" ; "empty text")]
|
|
fn test_extract_text_from_text_resource(input: &str, expected: &str) {
|
|
let resource = ResourceContents::TextResourceContents {
|
|
uri: "file:///test.txt".to_string(),
|
|
mime_type: Some("text/plain".to_string()),
|
|
text: input.to_string(),
|
|
meta: None,
|
|
};
|
|
assert_eq!(extract_text_from_resource(&resource), expected);
|
|
}
|
|
|
|
#[test_case("Hello from GitHub!", "Hello from GitHub!" ; "utf8 markdown")]
|
|
#[test_case("Simple text", "Simple text" ; "utf8 plain")]
|
|
fn test_extract_text_from_blob_utf8(input: &str, expected: &str) {
|
|
let blob = base64::engine::general_purpose::STANDARD.encode(input.as_bytes());
|
|
let resource = ResourceContents::BlobResourceContents {
|
|
uri: "github://repo/file.md".to_string(),
|
|
mime_type: Some("text/markdown".to_string()),
|
|
blob,
|
|
meta: None,
|
|
};
|
|
assert_eq!(extract_text_from_resource(&resource), expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_text_from_blob_binary() {
|
|
let binary_data: Vec<u8> = vec![0xFF, 0xFE, 0x00, 0x01, 0x89, 0x50, 0x4E, 0x47];
|
|
let blob = base64::engine::general_purpose::STANDARD.encode(&binary_data);
|
|
|
|
let resource = ResourceContents::BlobResourceContents {
|
|
uri: "file:///image.png".to_string(),
|
|
mime_type: Some("image/png".to_string()),
|
|
blob,
|
|
meta: None,
|
|
};
|
|
|
|
assert_eq!(
|
|
extract_text_from_resource(&resource),
|
|
"[Binary content (image/png) - 8 bytes]"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_text_from_blob_binary_no_mime_type() {
|
|
let binary_data: Vec<u8> = vec![0xFF, 0xFE];
|
|
let blob = base64::engine::general_purpose::STANDARD.encode(&binary_data);
|
|
|
|
let resource = ResourceContents::BlobResourceContents {
|
|
uri: "file:///unknown".to_string(),
|
|
mime_type: None,
|
|
blob,
|
|
meta: None,
|
|
};
|
|
|
|
assert_eq!(
|
|
extract_text_from_resource(&resource),
|
|
"[Binary content (application/octet-stream) - 2 bytes]"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_text_from_blob_invalid_base64() {
|
|
let resource = ResourceContents::BlobResourceContents {
|
|
uri: "file:///test.txt".to_string(),
|
|
mime_type: Some("text/plain".to_string()),
|
|
blob: "not valid base64!!!".to_string(),
|
|
meta: None,
|
|
};
|
|
assert_eq!(extract_text_from_resource(&resource), "not valid base64!!!");
|
|
}
|
|
}
|