Skip to main content

soma_provider_core/registry/
dispatch.rs

1use serde_json::Value;
2
3use crate::{ProviderCall, ProviderError, ProviderOutput};
4
5use super::{ProviderRegistry, RegisteredTool};
6
7impl ProviderRegistry {
8    pub async fn dispatch(&self, call: ProviderCall) -> Result<ProviderOutput, ProviderError> {
9        self.dispatch_with(
10            call,
11            |provider, call| async move { provider.call(call).await },
12        )
13        .await
14    }
15
16    pub async fn dispatch_with<F, Fut>(
17        &self,
18        call: ProviderCall,
19        invoke: F,
20    ) -> Result<ProviderOutput, ProviderError>
21    where
22        F: FnOnce(std::sync::Arc<dyn crate::Provider>, ProviderCall) -> Fut,
23        Fut: std::future::Future<Output = Result<ProviderOutput, ProviderError>>,
24    {
25        self.dispatch_with_pre_input(call, |_| Ok(()), invoke).await
26    }
27
28    /// Dispatches with a host check after action/surface resolution and before
29    /// provider-core validates the declared input limit and schema.
30    pub async fn dispatch_with_pre_input<P, F, Fut>(
31        &self,
32        mut call: ProviderCall,
33        pre_input: P,
34        invoke: F,
35    ) -> Result<ProviderOutput, ProviderError>
36    where
37        P: FnOnce(&ProviderCall) -> Result<(), ProviderError>,
38        F: FnOnce(std::sync::Arc<dyn crate::Provider>, ProviderCall) -> Fut,
39        Fut: std::future::Future<Output = Result<ProviderOutput, ProviderError>>,
40    {
41        let entry = self
42            .snapshot
43            .tool(&call.action)
44            .cloned()
45            .ok_or_else(|| ProviderError::tool_not_found(&call.action))?;
46        let provider = self
47            .providers
48            .get(entry.provider_id().as_str())
49            .cloned()
50            .ok_or_else(|| {
51                ProviderError::new(
52                    "provider_not_loaded",
53                    entry.provider_id().to_string(),
54                    Some(call.action.clone()),
55                    "provider is not loaded in the active registry",
56                    "Rebuild the provider registry and retry.",
57                )
58            })?;
59
60        call.provider = entry.provider_id().to_string();
61        call.snapshot_id = self.snapshot.fingerprint().to_string();
62        validate_surface(&entry, &call)?;
63        pre_input(&call)?;
64        validate_input(&entry, &call)?;
65        let output = invoke(provider, call).await?;
66        validate_output(&entry, &output)?;
67        Ok(output)
68    }
69}
70
71fn validate_surface(entry: &RegisteredTool, call: &ProviderCall) -> Result<(), ProviderError> {
72    if entry.spec().exposed_on(call.surface) {
73        return Ok(());
74    }
75    Err(ProviderError::validation(
76        entry.provider_id().to_string(),
77        &call.action,
78        "surface_not_exposed",
79        format!(
80            "action `{}` is not exposed on {}",
81            call.action,
82            call.surface.as_str()
83        ),
84    ))
85}
86
87fn validate_input(entry: &RegisteredTool, call: &ProviderCall) -> Result<(), ProviderError> {
88    if let Some(max) = entry
89        .spec()
90        .limits
91        .as_ref()
92        .and_then(|limits| limits.max_input_bytes)
93    {
94        let len = serialized_len(&call.params);
95        if len > max {
96            return Err(ProviderError::validation(
97                entry.provider_id().to_string(),
98                &call.action,
99                "input_too_large",
100                format!("provider input exceeded {max} bytes"),
101            ));
102        }
103    }
104    let details = entry
105        .input_validator
106        .iter_errors(&call.params)
107        .map(|error| format!("{}: {error}", error.instance_path()))
108        .collect::<Vec<_>>();
109    if details.is_empty() {
110        Ok(())
111    } else {
112        Err(ProviderError::validation(
113            entry.provider_id().to_string(),
114            &call.action,
115            "input_schema_failed",
116            details.join("; "),
117        ))
118    }
119}
120
121fn validate_output(entry: &RegisteredTool, output: &ProviderOutput) -> Result<(), ProviderError> {
122    if let Some(max) = entry
123        .spec()
124        .limits
125        .as_ref()
126        .and_then(|limits| limits.max_response_bytes)
127        && serialized_len(&output.value) > max
128    {
129        return Err(ProviderError::new(
130            "response_too_large",
131            entry.provider_id().to_string(),
132            Some(entry.spec().name.clone()),
133            format!("provider response exceeded {max} bytes"),
134            "Reduce the response size or add paging.",
135        ));
136    }
137    let Some(validator) = &entry.output_validator else {
138        return Ok(());
139    };
140    let details = validator
141        .iter_errors(&output.value)
142        .map(|error| format!("{}: {error}", error.instance_path()))
143        .collect::<Vec<_>>();
144    if details.is_empty() {
145        Ok(())
146    } else {
147        Err(ProviderError::new(
148            "output_schema_failed",
149            entry.provider_id().to_string(),
150            Some(entry.spec().name.clone()),
151            details.join("; "),
152            "Fix the provider output or its declared output schema, then retry.",
153        )
154        .with_phase("output_validation"))
155    }
156}
157
158fn serialized_len(value: &Value) -> usize {
159    serde_json::to_vec(value)
160        .map(|bytes| bytes.len())
161        .unwrap_or(usize::MAX)
162}