use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
#[cfg(test)]
use async_trait::async_trait;
use parking_lot::RwLock;
use crate::ports::{InternalUrlRouter, ProtocolHandler, ResolveContext, ResolvedUrl, SdkError};
pub struct CompositeUrlRouter {
handlers: RwLock<HashMap<String, Arc<dyn ProtocolHandler>>>,
}
impl CompositeUrlRouter {
pub fn new() -> Self {
Self {
handlers: RwLock::new(HashMap::new()),
}
}
pub fn register(&self, handler: Arc<dyn ProtocolHandler>) {
self.handlers
.write()
.insert(handler.scheme().to_lowercase(), handler);
}
pub fn unregister(&self, scheme: &str) -> bool {
self.handlers
.write()
.remove(&scheme.to_lowercase())
.is_some()
}
pub fn can_resolve(&self, scheme: &str) -> bool {
self.handlers.read().contains_key(&scheme.to_lowercase())
}
fn parse_scheme(input: &str) -> Option<&str> {
let bytes = input.as_bytes();
let len = bytes.len();
if len < 3 {
return None;
}
if !bytes[0].is_ascii_lowercase() {
return None;
}
for i in 1..len {
match bytes[i] {
b':' => {
if i + 2 < len && bytes[i + 1] == b'/' && bytes[i + 2] == b'/' {
let scheme = &input[..i];
if scheme.eq_ignore_ascii_case("http")
|| scheme.eq_ignore_ascii_case("https")
{
return None;
}
return Some(scheme);
}
return None;
}
b if b.is_ascii_lowercase()
|| b.is_ascii_digit()
|| b == b'+'
|| b == b'.'
|| b == b'-' => {}
_ => return None,
}
}
None
}
fn split_uri(uri: &str) -> Option<(&str, &str)> {
let scheme = Self::parse_scheme(uri)?;
let path = &uri[scheme.len() + 3..]; Some((scheme, path))
}
}
impl Default for CompositeUrlRouter {
fn default() -> Self {
Self::new()
}
}
impl InternalUrlRouter for CompositeUrlRouter {
fn schemes(&self) -> &[&str] {
&[]
}
fn registered_schemes(&self) -> Vec<String> {
self.handlers.read().keys().cloned().collect()
}
fn resolve<'a>(
&'a self,
uri: &'a str,
ctx: &'a ResolveContext,
) -> Pin<Box<dyn Future<Output = Result<ResolvedUrl, SdkError>> + Send + 'a>> {
Box::pin(async move {
let (scheme, path) = Self::split_uri(uri).ok_or_else(|| SdkError::UnknownScheme {
scheme: uri.to_string(),
})?;
let handler = self.handlers.read().get(scheme).cloned().ok_or_else(|| {
SdkError::UnknownScheme {
scheme: scheme.to_string(),
}
})?;
let mut resolved = handler.resolve(path, None, ctx).await?;
resolved.immutable = resolved.immutable || handler.immutable();
Ok(resolved)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_scheme_valid() {
assert_eq!(
CompositeUrlRouter::parse_scheme("issue://1428"),
Some("issue")
);
assert_eq!(
CompositeUrlRouter::parse_scheme("pr://owner/repo/1428"),
Some("pr")
);
assert_eq!(
CompositeUrlRouter::parse_scheme("agent://sub1/output"),
Some("agent")
);
assert_eq!(
CompositeUrlRouter::parse_scheme("skill://my-skill/SKILL.md"),
Some("skill")
);
assert_eq!(
CompositeUrlRouter::parse_scheme("memory://session-abc"),
Some("memory")
);
}
#[test]
fn test_parse_scheme_rejects_regular_paths() {
assert_eq!(CompositeUrlRouter::parse_scheme("src/main.rs"), None);
assert_eq!(CompositeUrlRouter::parse_scheme("/absolute/path"), None,);
assert_eq!(CompositeUrlRouter::parse_scheme("relative/path"), None);
}
#[test]
fn test_parse_scheme_rejects_http() {
assert_eq!(
CompositeUrlRouter::parse_scheme("https://example.com"),
None,
);
}
#[test]
fn test_split_uri() {
let (scheme, path) = CompositeUrlRouter::split_uri("pr://owner/repo/1428/diff/1").unwrap();
assert_eq!(scheme, "pr");
assert_eq!(path, "owner/repo/1428/diff/1");
}
#[test]
fn test_can_resolve() {
let router = CompositeUrlRouter::new();
assert!(!router.can_resolve("issue"));
struct DummyHandler;
#[async_trait]
impl ProtocolHandler for DummyHandler {
fn scheme(&self) -> &str {
"issue"
}
async fn resolve(
&self,
_: &str,
_: Option<&str>,
_: &ResolveContext,
) -> Result<ResolvedUrl, SdkError> {
Ok(ResolvedUrl {
url: "issue://1".into(),
content: "test".into(),
content_type: "text/plain".into(),
size: None,
source_path: None,
notes: vec![],
immutable: false,
})
}
}
router.register(Arc::new(DummyHandler));
assert!(router.can_resolve("issue"));
assert!(!router.can_resolve("pr"));
}
#[test]
fn test_unregister() {
let router = CompositeUrlRouter::new();
assert!(!router.unregister("issue"));
struct DummyHandler;
#[async_trait]
impl ProtocolHandler for DummyHandler {
fn scheme(&self) -> &str {
"issue"
}
async fn resolve(
&self,
_: &str,
_: Option<&str>,
_: &ResolveContext,
) -> Result<ResolvedUrl, SdkError> {
Ok(ResolvedUrl {
url: "issue://1".into(),
content: "test".into(),
content_type: "text/plain".into(),
size: None,
source_path: None,
notes: vec![],
immutable: false,
})
}
}
router.register(Arc::new(DummyHandler));
assert!(router.unregister("issue"));
assert!(!router.can_resolve("issue"));
}
}