use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use async_trait::async_trait;
use crate::{
ContentBlock, Extensions, UniversalItem, UniversalRequest, UniversalResponse, UniversalTool,
};
pub type ServiceSideResult<T> = std::result::Result<T, ServiceSideError>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ServiceSideError {
InvalidRegistration { message: String },
InvalidInput { capability: String, message: String },
Execution { capability: String, message: String },
}
impl ServiceSideError {
pub fn invalid_input(capability: impl Into<String>, message: impl Into<String>) -> Self {
Self::InvalidInput {
capability: capability.into(),
message: message.into(),
}
}
pub fn execution(capability: impl Into<String>, message: impl Into<String>) -> Self {
Self::Execution {
capability: capability.into(),
message: message.into(),
}
}
fn invalid_registration(message: impl Into<String>) -> Self {
Self::InvalidRegistration {
message: message.into(),
}
}
}
impl fmt::Display for ServiceSideError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidRegistration { message } => {
write!(f, "invalid service-side capability registration: {message}")
}
Self::InvalidInput {
capability,
message,
} => write!(
f,
"service-side capability '{capability}' rejected input: {message}"
),
Self::Execution {
capability,
message,
} => write!(
f,
"service-side capability '{capability}' failed: {message}"
),
}
}
}
impl std::error::Error for ServiceSideError {}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum ServiceSideInputKind {
Image,
File,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ServiceSideInput {
pub kind: ServiceSideInputKind,
pub media_type: Option<String>,
pub filename: Option<String>,
pub url: Option<String>,
pub data: Option<String>,
pub context: String,
pub extensions: Extensions,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ServiceSideInputResolution {
pub text: String,
}
#[async_trait]
pub trait ServiceSideInputResolver: Send + Sync {
fn id(&self) -> &str;
fn input_kind(&self) -> ServiceSideInputKind;
async fn resolve(
&self,
input: ServiceSideInput,
) -> ServiceSideResult<ServiceSideInputResolution>;
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ServiceSideResolutionReport {
pub images_resolved: usize,
pub files_resolved: usize,
}
impl ServiceSideResolutionReport {
pub fn changed(&self) -> bool {
self.images_resolved > 0 || self.files_resolved > 0
}
fn record(&mut self, kind: ServiceSideInputKind) {
match kind {
ServiceSideInputKind::Image => self.images_resolved += 1,
ServiceSideInputKind::File => self.files_resolved += 1,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ServiceSideToolCall {
pub id: String,
pub name: String,
pub arguments: serde_json::Value,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ServiceSideToolOutput {
pub content: Vec<ContentBlock>,
pub is_error: bool,
}
#[async_trait]
pub trait ServiceSideTool: Send + Sync {
fn definition(&self) -> UniversalTool;
async fn call(&self, call: ServiceSideToolCall) -> ServiceSideResult<ServiceSideToolOutput>;
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ServiceSideToolExecution {
pub results: Vec<UniversalItem>,
pub unhandled_calls: Vec<ServiceSideToolCall>,
}
impl ServiceSideToolExecution {
pub fn handled_count(&self) -> usize {
self.results.len()
}
pub fn append_results_to(self, request: &mut UniversalRequest) {
request.input.extend(self.results);
}
}
struct RegisteredTool {
definition: UniversalTool,
handler: Arc<dyn ServiceSideTool>,
}
#[derive(Default)]
pub struct ServiceSideCapabilityRegistry {
input_resolvers: BTreeMap<ServiceSideInputKind, Arc<dyn ServiceSideInputResolver>>,
tools: BTreeMap<String, RegisteredTool>,
}
impl ServiceSideCapabilityRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register_input_resolver<R>(&mut self, resolver: R) -> ServiceSideResult<()>
where
R: ServiceSideInputResolver + 'static,
{
let id = resolver.id().trim();
if id.is_empty() {
return Err(ServiceSideError::invalid_registration(
"input resolver id cannot be empty",
));
}
let kind = resolver.input_kind();
if let Some(existing) = self.input_resolvers.get(&kind) {
return Err(ServiceSideError::invalid_registration(format!(
"input kind {kind:?} is already owned by resolver '{}'",
existing.id()
)));
}
self.input_resolvers.insert(kind, Arc::new(resolver));
Ok(())
}
pub fn register_tool<T>(&mut self, tool: T) -> ServiceSideResult<()>
where
T: ServiceSideTool + 'static,
{
let definition = tool.definition();
let name = definition.name.trim();
if name.is_empty() {
return Err(ServiceSideError::invalid_registration(
"host tool name cannot be empty",
));
}
if self.tools.contains_key(name) {
return Err(ServiceSideError::invalid_registration(format!(
"host tool '{name}' is already registered"
)));
}
self.tools.insert(
name.to_string(),
RegisteredTool {
definition,
handler: Arc::new(tool),
},
);
Ok(())
}
pub fn inject_tools(&self, request: &mut UniversalRequest) -> usize {
let existing = request
.tools
.iter()
.map(|tool| tool.name.as_str())
.collect::<BTreeSet<_>>();
let additions = self
.tools
.values()
.filter(|tool| !existing.contains(tool.definition.name.as_str()))
.map(|tool| tool.definition.clone())
.collect::<Vec<_>>();
let count = additions.len();
request.tools.extend(additions);
count
}
pub async fn resolve_inputs(
&self,
request: &mut UniversalRequest,
) -> ServiceSideResult<ServiceSideResolutionReport> {
let mut report = ServiceSideResolutionReport::default();
self.resolve_blocks(&mut request.instructions, &mut report)
.await?;
for item in &mut request.input {
match item {
UniversalItem::Message { content, .. }
| UniversalItem::ToolResult { content, .. } => {
self.resolve_blocks(content, &mut report).await?;
}
_ => {}
}
}
Ok(report)
}
pub async fn execute_tool_calls(
&self,
response: &UniversalResponse,
) -> ServiceSideResult<ServiceSideToolExecution> {
let mut execution = ServiceSideToolExecution::default();
for call in collect_tool_calls(response) {
let Some(tool) = self.tools.get(&call.name) else {
execution.unhandled_calls.push(call);
continue;
};
let output = tool.handler.call(call.clone()).await?;
execution.results.push(UniversalItem::ToolResult {
tool_call_id: call.id,
content: output.content,
is_error: output.is_error,
extensions: Extensions::default(),
});
}
Ok(execution)
}
fn resolve_blocks<'a>(
&'a self,
blocks: &'a mut [ContentBlock],
report: &'a mut ServiceSideResolutionReport,
) -> Pin<Box<dyn Future<Output = ServiceSideResult<()>> + Send + 'a>> {
Box::pin(async move {
let context = adjacent_text(blocks);
for block in blocks {
let input = match block {
ContentBlock::Image {
media_type,
url,
data,
extensions,
} => Some(ServiceSideInput {
kind: ServiceSideInputKind::Image,
media_type: media_type.clone(),
filename: None,
url: url.clone(),
data: data.clone(),
context: context.clone(),
extensions: extensions.clone(),
}),
ContentBlock::File {
media_type,
filename,
url,
data,
extensions,
} => Some(ServiceSideInput {
kind: ServiceSideInputKind::File,
media_type: media_type.clone(),
filename: filename.clone(),
url: url.clone(),
data: data.clone(),
context: context.clone(),
extensions: extensions.clone(),
}),
ContentBlock::ToolResult { content, .. } => {
self.resolve_blocks(content, report).await?;
None
}
_ => None,
};
let Some(input) = input else {
continue;
};
let Some(resolver) = self.input_resolvers.get(&input.kind) else {
continue;
};
let kind = input.kind;
let resolution = resolver.resolve(input).await?;
if resolution.text.trim().is_empty() {
return Err(ServiceSideError::Execution {
capability: resolver.id().to_string(),
message: "input resolver returned empty replacement text".to_string(),
});
}
*block = ContentBlock::Text {
text: resolution.text,
};
report.record(kind);
}
Ok(())
})
}
}
pub fn request_contains_service_side_input(
request: &UniversalRequest,
kind: ServiceSideInputKind,
) -> bool {
blocks_contain_input(&request.instructions, kind)
|| request.input.iter().any(|item| match item {
UniversalItem::Message { content, .. } | UniversalItem::ToolResult { content, .. } => {
blocks_contain_input(content, kind)
}
_ => false,
})
}
fn blocks_contain_input(blocks: &[ContentBlock], kind: ServiceSideInputKind) -> bool {
blocks.iter().any(|block| match block {
ContentBlock::Image { .. } => kind == ServiceSideInputKind::Image,
ContentBlock::File { .. } => kind == ServiceSideInputKind::File,
ContentBlock::ToolResult { content, .. } => blocks_contain_input(content, kind),
_ => false,
})
}
fn adjacent_text(blocks: &[ContentBlock]) -> String {
blocks
.iter()
.filter_map(|block| match block {
ContentBlock::Text { text } => Some(text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("\n")
}
fn collect_tool_calls(response: &UniversalResponse) -> Vec<ServiceSideToolCall> {
let mut calls = Vec::new();
let mut ids = BTreeSet::new();
for item in &response.output {
match item {
UniversalItem::ToolCall {
id,
name,
arguments,
..
} => push_tool_call(&mut calls, &mut ids, id, name, arguments),
UniversalItem::Message { content, .. } | UniversalItem::ToolResult { content, .. } => {
collect_block_tool_calls(content, &mut calls, &mut ids)
}
_ => {}
}
}
calls
}
fn collect_block_tool_calls(
blocks: &[ContentBlock],
calls: &mut Vec<ServiceSideToolCall>,
ids: &mut BTreeSet<String>,
) {
for block in blocks {
match block {
ContentBlock::ToolCall {
id,
name,
arguments,
..
} => push_tool_call(calls, ids, id, name, arguments),
ContentBlock::ToolResult { content, .. } => {
collect_block_tool_calls(content, calls, ids)
}
_ => {}
}
}
}
fn push_tool_call(
calls: &mut Vec<ServiceSideToolCall>,
ids: &mut BTreeSet<String>,
id: &str,
name: &str,
arguments: &serde_json::Value,
) {
if ids.insert(id.to_string()) {
calls.push(ServiceSideToolCall {
id: id.to_string(),
name: name.to_string(),
arguments: arguments.clone(),
});
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use serde_json::json;
use super::*;
use crate::Role;
struct ImageResolver;
#[async_trait]
impl ServiceSideInputResolver for ImageResolver {
fn id(&self) -> &str {
"test_image_resolver"
}
fn input_kind(&self) -> ServiceSideInputKind {
ServiceSideInputKind::Image
}
async fn resolve(
&self,
input: ServiceSideInput,
) -> ServiceSideResult<ServiceSideInputResolution> {
assert_eq!(input.context, "What is visible?");
assert_eq!(input.data.as_deref(), Some("aW1hZ2U="));
Ok(ServiceSideInputResolution {
text: "A terminal window.".to_string(),
})
}
}
#[test]
fn rejects_duplicate_input_resolvers_for_one_kind() {
let mut registry = ServiceSideCapabilityRegistry::new();
registry
.register_input_resolver(ImageResolver)
.expect("first resolver registers");
let error = registry
.register_input_resolver(ImageResolver)
.expect_err("second image resolver is ambiguous");
assert!(error.to_string().contains("already owned"));
}
#[test]
fn registered_resolver_replaces_media_in_place() {
let mut registry = ServiceSideCapabilityRegistry::new();
registry
.register_input_resolver(ImageResolver)
.expect("resolver registers");
let mut request = UniversalRequest {
input: vec![UniversalItem::Message {
role: Role::User,
id: None,
content: vec![
ContentBlock::Text {
text: "What is visible?".to_string(),
},
ContentBlock::Image {
media_type: Some("image/png".to_string()),
url: None,
data: Some("aW1hZ2U=".to_string()),
extensions: Extensions::default(),
},
],
extensions: Extensions::default(),
}],
..UniversalRequest::default()
};
let report = futures_lite::future::block_on(registry.resolve_inputs(&mut request))
.expect("image resolves");
assert_eq!(report.images_resolved, 1);
let UniversalItem::Message { content, .. } = &request.input[0] else {
panic!("message remains a message");
};
assert_eq!(
content[1],
ContentBlock::Text {
text: "A terminal window.".to_string()
}
);
}
#[test]
fn detects_registered_input_kind_inside_nested_tool_result() {
let request = UniversalRequest {
input: vec![UniversalItem::ToolResult {
tool_call_id: "call-image".to_string(),
content: vec![ContentBlock::ToolResult {
tool_call_id: "nested-image".to_string(),
content: vec![ContentBlock::Image {
media_type: Some("image/png".to_string()),
url: None,
data: Some("aW1hZ2U=".to_string()),
extensions: Extensions::default(),
}],
is_error: false,
extensions: Extensions::default(),
}],
is_error: false,
extensions: Extensions::default(),
}],
..UniversalRequest::default()
};
assert!(request_contains_service_side_input(
&request,
ServiceSideInputKind::Image
));
assert!(!request_contains_service_side_input(
&request,
ServiceSideInputKind::File
));
}
struct EchoTool {
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl ServiceSideTool for EchoTool {
fn definition(&self) -> UniversalTool {
UniversalTool {
name: "host_echo".to_string(),
description: Some("Echo host-side text".to_string()),
input_schema: Some(json!({
"type": "object",
"properties": { "text": { "type": "string" } }
})),
strict: None,
extensions: Extensions::default(),
}
}
async fn call(
&self,
call: ServiceSideToolCall,
) -> ServiceSideResult<ServiceSideToolOutput> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(ServiceSideToolOutput {
content: vec![ContentBlock::Text {
text: call.arguments["text"]
.as_str()
.unwrap_or_default()
.to_string(),
}],
is_error: false,
})
}
}
#[test]
fn registry_injects_and_executes_registered_tools_only() {
let calls = Arc::new(AtomicUsize::new(0));
let mut registry = ServiceSideCapabilityRegistry::new();
registry
.register_tool(EchoTool {
calls: calls.clone(),
})
.expect("tool registers");
let mut request = UniversalRequest::default();
assert_eq!(registry.inject_tools(&mut request), 1);
assert_eq!(registry.inject_tools(&mut request), 0);
let response = UniversalResponse {
output: vec![
UniversalItem::ToolCall {
id: "call-host".to_string(),
name: "host_echo".to_string(),
arguments: json!({ "text": "hello" }),
extensions: Extensions::default(),
},
UniversalItem::ToolCall {
id: "call-client".to_string(),
name: "client_tool".to_string(),
arguments: json!({}),
extensions: Extensions::default(),
},
],
..UniversalResponse::default()
};
let execution = futures_lite::future::block_on(registry.execute_tool_calls(&response))
.expect("registered tool executes");
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert_eq!(execution.handled_count(), 1);
assert_eq!(execution.unhandled_calls.len(), 1);
let UniversalItem::ToolResult { content, .. } = &execution.results[0] else {
panic!("host call becomes tool result");
};
assert_eq!(
content[0],
ContentBlock::Text {
text: "hello".to_string()
}
);
}
}