va_ai_api_bridge/

service_side.rs

//! Host-registered service-side capabilities.
//!
//! The registry owns orchestration over [`UniversalRequest`] and
//! [`UniversalResponse`]. Embedding hosts still own every external effect:
//! HTTP, credentials, caching, persistence, retries, and sandbox policy.

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>,
    /// Text blocks adjacent to this input in the same message or tool result.
    pub context: String,
    pub extensions: Extensions,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ServiceSideInputResolution {
    /// Replacement text inserted at the original input position.
    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(())
    }

    /// Inject every registered host tool into the request, preserving any
    /// client declaration with the same name as the request owner.
    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()
            }
        );
    }
}