rpi_runtime/provider/
factories.rs

1//! Built-in protocols and native extensions implement the same factory contract.
2//! Authentication and model binding are applied once, outside every factory.
3use super::binding::{ExecutionBinding, RequestAuth};
4use super::*;
5use crate::ProviderDeclaration as ProviderManifest;
6
7trait ProviderFactory: Send + Sync {
8    fn create(
9        &self,
10        provider_id: &str,
11        models: Vec<Model>,
12        settings: &crate::RuntimeSettings,
13        values: &dyn ValueSource,
14    ) -> Result<Arc<dyn Provider>, ResolveError>;
15}
16
17struct ProtocolFactory(rpi_ai::Api);
18impl ProviderFactory for ProtocolFactory {
19    fn create(
20        &self,
21        provider_id: &str,
22        models: Vec<Model>,
23        settings: &crate::RuntimeSettings,
24        values: &dyn ValueSource,
25    ) -> Result<Arc<dyn Provider>, ResolveError> {
26        let fallback = if self.0 == rpi_ai::Api::AnthropicMessages {
27            "https://api.anthropic.com"
28        } else {
29            "https://api.openai.com/v1"
30        };
31        let http = provider_http_client(&models, fallback, settings, values)?;
32        let provider: Arc<dyn Provider> = match self.0 {
33            rpi_ai::Api::AnthropicMessages => Arc::new(
34                AnthropicProvider::with_models_without_env_api_key(None, http, models),
35            ),
36            rpi_ai::Api::OpenaiCompletions => {
37                Arc::new(OpenAiCompletionsProvider::with_models_without_env_api_key(
38                    provider_id,
39                    None,
40                    http,
41                    models,
42                ))
43            }
44            rpi_ai::Api::OpenaiResponses => {
45                Arc::new(OpenAiResponsesProvider::with_models_without_env_api_key(
46                    provider_id,
47                    None,
48                    http,
49                    models,
50                ))
51            }
52            _ => unreachable!("only supported wire protocols are registered"),
53        };
54        Ok(provider)
55    }
56}
57
58struct ExtensionFactory(Arc<dyn Provider>);
59impl ProviderFactory for ExtensionFactory {
60    fn create(
61        &self,
62        _: &str,
63        _: Vec<Model>,
64        _: &crate::RuntimeSettings,
65        _: &dyn ValueSource,
66    ) -> Result<Arc<dyn Provider>, ResolveError> {
67        Ok(self.0.clone())
68    }
69}
70
71pub struct RuntimeFactories {
72    protocols: BTreeMap<String, Arc<dyn ProviderFactory>>,
73    providers: BTreeMap<(String, String), Arc<dyn ProviderFactory>>,
74    pub manifests: Vec<ProviderManifest>,
75}
76
77impl Default for RuntimeFactories {
78    fn default() -> Self {
79        let protocols = [
80            rpi_ai::Api::AnthropicMessages,
81            rpi_ai::Api::OpenaiCompletions,
82            rpi_ai::Api::OpenaiResponses,
83        ]
84        .into_iter()
85        .map(|api| {
86            (
87                api.as_str().to_string(),
88                Arc::new(ProtocolFactory(api)) as Arc<dyn ProviderFactory>,
89            )
90        })
91        .collect();
92        Self {
93            protocols,
94            providers: BTreeMap::new(),
95            manifests: Vec::new(),
96        }
97    }
98}
99
100impl RuntimeFactories {
101    /// Add a declaration without loading an implementation (catalog discovery).
102    pub fn declare(&mut self, declaration: ProviderManifest) {
103        self.manifests.retain(|known| known.id != declaration.id);
104        self.manifests.push(declaration);
105    }
106    /// Register a host-owned implementation. Its Arc keeps the original loader
107    /// lease alive for all requests resolved from this registry.
108    pub fn register(&mut self, provider: Arc<dyn Provider>, api: rpi_ai::Api) {
109        let id = provider.id().to_string();
110        self.providers.insert(
111            (id, api.as_str().into()),
112            Arc::new(ExtensionFactory(provider)),
113        );
114    }
115
116    pub(super) fn configure(
117        &self,
118        provider_id: &str,
119        api: &rpi_ai::Api,
120        models: Vec<Model>,
121        auth: RequestAuth,
122        settings: &crate::RuntimeSettings,
123        values: &dyn ValueSource,
124    ) -> Result<(Arc<dyn Provider>, Arc<ExecutionBinding>), ResolveError> {
125        let factory = self
126            .providers
127            .get(&(provider_id.into(), api.as_str().into()))
128            .or_else(|| self.protocols.get(api.as_str()))
129            .ok_or_else(|| ResolveError::UnsupportedApi {
130                provider: provider_id.into(),
131                api: api.as_str().into(),
132            })?;
133        let binding = Arc::new(ExecutionBinding::new(provider_id, models, auth));
134        let inner = factory.create(
135            provider_id,
136            binding.descriptors().to_vec(),
137            settings,
138            values,
139        )?;
140        let provider = Arc::new(ConfiguredProvider {
141            id: provider_id.into(),
142            inner,
143            binding: binding.clone(),
144        });
145        Ok((provider, binding))
146    }
147}
148
149struct ConfiguredProvider {
150    id: String,
151    inner: Arc<dyn Provider>,
152    binding: Arc<ExecutionBinding>,
153}
154
155/// The harness resolves providers by ID. Keep each ID unique even when a
156/// configured account and an extension expose different wire protocols.
157pub(super) fn combine_provider_runtimes(
158    providers: Vec<Arc<dyn Provider>>,
159) -> Vec<Arc<dyn Provider>> {
160    let mut groups: Vec<Vec<Arc<dyn Provider>>> = Vec::new();
161    for provider in providers {
162        if let Some(group) = groups
163            .iter_mut()
164            .find(|group| group[0].id() == provider.id())
165        {
166            group.push(provider);
167        } else {
168            groups.push(vec![provider]);
169        }
170    }
171    groups
172        .into_iter()
173        .map(|mut routes| {
174            if routes.len() == 1 {
175                return routes.pop().unwrap();
176            }
177            Arc::new(RoutedProvider {
178                id: routes[0].id().into(),
179                models: routes
180                    .iter()
181                    .flat_map(|route| route.models().to_vec())
182                    .collect(),
183                routes,
184            }) as Arc<dyn Provider>
185        })
186        .collect()
187}
188
189struct RoutedProvider {
190    id: String,
191    models: Vec<Model>,
192    routes: Vec<Arc<dyn Provider>>,
193}
194
195#[async_trait::async_trait]
196impl Provider for RoutedProvider {
197    fn id(&self) -> &str {
198        &self.id
199    }
200    fn models(&self) -> &[Model] {
201        &self.models
202    }
203    async fn stream_simple(
204        &self,
205        model: &Model,
206        context: &Context,
207        options: &SimpleStreamOptions,
208    ) -> AssistantMessageEventStream {
209        let route = self.routes.iter().find(|route| {
210            route.models().iter().any(|candidate| {
211                candidate.provider == model.provider
212                    && candidate.api == model.api
213                    && candidate.id.eq_ignore_ascii_case(&model.id)
214            })
215        });
216        match route {
217            Some(route) => route.stream_simple(model, context, options).await,
218            None => unbound_model_error(model, &self.id).await,
219        }
220    }
221}
222
223async fn unbound_model_error(model: &Model, provider_id: &str) -> AssistantMessageEventStream {
224    let (mut producer, stream) = rpi_ai::event_stream::create_assistant_message_event_stream();
225    producer
226        .push(rpi_ai::AssistantMessageEvent::Error {
227            reason: rpi_ai::types::ErrorReason::Error,
228            error: rpi_ai::types::AssistantMessage::terminal(
229                model.api.clone(),
230                provider_id,
231                model.id.clone(),
232                rpi_ai::types::StopReason::Error,
233                "model is not registered in this provider execution binding",
234                0,
235            ),
236        })
237        .await;
238    stream
239}
240
241#[async_trait::async_trait]
242impl Provider for ConfiguredProvider {
243    fn id(&self) -> &str {
244        &self.id
245    }
246    fn models(&self) -> &[Model] {
247        self.binding.descriptors()
248    }
249    async fn stream_simple(
250        &self,
251        model: &Model,
252        context: &Context,
253        options: &SimpleStreamOptions,
254    ) -> AssistantMessageEventStream {
255        let Some((model, options)) = self.binding.prepare(model, options) else {
256            return unbound_model_error(model, &self.id).await;
257        };
258        self.inner.stream_simple(&model, context, &options).await
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use rpi_ai::{
266        types::{AssistantMessage, DoneReason, ErrorReason, StopReason},
267        Api, AssistantMessageEvent,
268    };
269    use std::sync::Mutex;
270
271    type Calls = Arc<Mutex<Vec<(Model, SimpleStreamOptions)>>>;
272    struct RecordingFactory(Calls);
273    struct RecordingProvider(Calls);
274    impl ProviderFactory for RecordingFactory {
275        fn create(
276            &self,
277            _: &str,
278            models: Vec<Model>,
279            _: &crate::RuntimeSettings,
280            _: &dyn ValueSource,
281        ) -> Result<Arc<dyn Provider>, ResolveError> {
282            assert!(
283                models.iter().all(|model| model.headers.is_none()),
284                "factories receive public descriptors"
285            );
286            Ok(Arc::new(RecordingProvider(self.0.clone())))
287        }
288    }
289    #[async_trait::async_trait]
290    impl Provider for RecordingProvider {
291        fn id(&self) -> &str {
292            "recording"
293        }
294        fn models(&self) -> &[Model] {
295            &[]
296        }
297        async fn stream_simple(
298            &self,
299            model: &Model,
300            _: &Context,
301            options: &SimpleStreamOptions,
302        ) -> AssistantMessageEventStream {
303            self.0
304                .lock()
305                .unwrap()
306                .push((model.clone(), options.clone()));
307            let (mut producer, stream) =
308                rpi_ai::event_stream::create_assistant_message_event_stream();
309            producer
310                .push(AssistantMessageEvent::Done {
311                    reason: DoneReason::Stop,
312                    message: AssistantMessage::empty(
313                        model.api.clone(),
314                        model.provider.clone(),
315                        model.id.clone(),
316                        0,
317                    ),
318                })
319                .await;
320            stream
321        }
322    }
323
324    #[tokio::test]
325    async fn one_provider_routes_duplicate_model_ids_by_api_with_private_credentials() {
326        let calls = Calls::default();
327        let mut factories = RuntimeFactories::default();
328        for api in [Api::OpenaiCompletions, Api::OpenaiResponses] {
329            factories.protocols.insert(
330                api.as_str().into(),
331                Arc::new(RecordingFactory(calls.clone())),
332            );
333        }
334        let mut runtimes = Vec::new();
335        for (api, token) in [
336            (Api::OpenaiCompletions, "completions-token"),
337            (Api::OpenaiResponses, "responses-token"),
338        ] {
339            let mut model = Model::new(
340                "same-model",
341                "Same model",
342                api.clone(),
343                "deepseek",
344                format!("https://configured.example/{}", api.as_str()),
345            );
346            model.headers = Some(BTreeMap::from([("x-route".into(), token.into())]));
347            runtimes.push(
348                factories
349                    .configure(
350                        "deepseek",
351                        &api,
352                        vec![model],
353                        RequestAuth::ApiKey(token.into()),
354                        &crate::RuntimeSettings::default(),
355                        &BTreeMap::<String, String>::new(),
356                    )
357                    .unwrap()
358                    .0,
359            );
360        }
361        let providers = combine_provider_runtimes(runtimes);
362        assert_eq!(providers.len(), 1);
363        let provider = &providers[0];
364        assert_eq!(provider.models().len(), 2);
365        assert!(provider
366            .models()
367            .iter()
368            .all(|model| model.headers.is_none()));
369        // This is the same ID-only lookup used by the harness after /provider.
370        for model in provider.models() {
371            let selected = providers
372                .iter()
373                .find(|provider| provider.id() == model.provider)
374                .unwrap();
375            let mut requested = model.clone();
376            requested.id = requested.id.to_uppercase();
377            requested.base_url = "https://untrusted.example".into();
378            selected
379                .stream_simple(
380                    &requested,
381                    &Context::new(vec![]),
382                    &SimpleStreamOptions::default(),
383                )
384                .await
385                .result()
386                .await
387                .unwrap();
388        }
389        let recorded = calls.lock().unwrap().clone();
390        assert_eq!(recorded.len(), 2);
391        for (call, (api, token)) in recorded.iter().zip([
392            (Api::OpenaiCompletions, "completions-token"),
393            (Api::OpenaiResponses, "responses-token"),
394        ]) {
395            assert_eq!(call.0.api, api);
396            assert_eq!(
397                call.0.base_url,
398                format!("https://configured.example/{}", api.as_str())
399            );
400            assert_eq!(call.0.headers.as_ref().unwrap()["x-route"], token);
401            assert_eq!(call.1.api_key.as_deref(), Some(token));
402        }
403        drop(recorded);
404        for (provider_id, api, id) in [
405            ("other", Api::OpenaiCompletions, "same-model"),
406            ("deepseek", Api::AnthropicMessages, "same-model"),
407            ("deepseek", Api::OpenaiResponses, "unknown-model"),
408        ] {
409            let requested = Model::new(id, id, api, provider_id, "https://untrusted.example");
410            let result = provider
411                .stream_simple(
412                    &requested,
413                    &Context::new(vec![]),
414                    &SimpleStreamOptions::default(),
415                )
416                .await
417                .result()
418                .await
419                .unwrap();
420            assert_eq!(result.stop_reason, StopReason::Error);
421        }
422        assert_eq!(calls.lock().unwrap().len(), 2);
423    }
424
425    #[tokio::test]
426    async fn extension_factory_only_handles_its_registered_wire_api() {
427        let extension_calls = Calls::default();
428        let protocol_calls = Calls::default();
429        let mut factories = RuntimeFactories::default();
430        factories.providers.insert(
431            ("deepseek".into(), "openai-completions".into()),
432            Arc::new(RecordingFactory(extension_calls.clone())),
433        );
434        factories.protocols.insert(
435            "openai-responses".into(),
436            Arc::new(RecordingFactory(protocol_calls.clone())),
437        );
438        for api in [Api::OpenaiCompletions, Api::OpenaiResponses] {
439            let model = Model::new(
440                "same-model",
441                "Same model",
442                api.clone(),
443                "deepseek",
444                "https://configured.example/v1",
445            );
446            let (provider, _) = factories
447                .configure(
448                    "deepseek",
449                    &api,
450                    vec![model],
451                    RequestAuth::None,
452                    &crate::RuntimeSettings::default(),
453                    &BTreeMap::<String, String>::new(),
454                )
455                .unwrap();
456            provider
457                .stream_simple(
458                    &provider.models()[0],
459                    &Context::new(vec![]),
460                    &SimpleStreamOptions::default(),
461                )
462                .await
463                .result()
464                .await
465                .unwrap();
466        }
467        assert_eq!(extension_calls.lock().unwrap().len(), 1);
468        assert_eq!(
469            extension_calls.lock().unwrap()[0].0.api,
470            Api::OpenaiCompletions
471        );
472        assert_eq!(protocol_calls.lock().unwrap().len(), 1);
473        assert_eq!(
474            protocol_calls.lock().unwrap()[0].0.api,
475            Api::OpenaiResponses
476        );
477    }
478
479    #[tokio::test]
480    async fn protocol_and_extension_factories_use_the_same_private_binding() {
481        for extension in [false, true] {
482            let calls = Calls::default();
483            let mut factories = RuntimeFactories::default();
484            let factory = Arc::new(RecordingFactory(calls.clone())) as Arc<dyn ProviderFactory>;
485            if extension {
486                factories
487                    .providers
488                    .insert(("account".into(), "openai-completions".into()), factory);
489            } else {
490                factories
491                    .protocols
492                    .insert("openai-completions".into(), factory);
493            }
494            let mut raw = Model::new(
495                "model",
496                "Model",
497                Api::OpenaiCompletions,
498                "account",
499                "https://configured.example/v1",
500            );
501            raw.headers = Some(BTreeMap::from([(
502                "x-api-key".into(),
503                "private-header".into(),
504            )]));
505            let (provider, binding) = factories
506                .configure(
507                    "account",
508                    &raw.api.clone(),
509                    vec![raw],
510                    RequestAuth::ApiKey("private-credential".into()),
511                    &crate::RuntimeSettings::default(),
512                    &BTreeMap::<String, String>::new(),
513                )
514                .unwrap();
515            let serialized = serde_json::to_string(provider.models()).unwrap();
516            assert!(
517                !serialized.contains("private-header")
518                    && !serialized.contains("private-credential")
519            );
520            assert!(!format!("{binding:?}").contains("private-credential"));
521            let mut requested = provider.models()[0].clone();
522            requested.base_url = "https://untrusted.example/v1".into();
523            requested.headers = Some(BTreeMap::from([(
524                "authorization".into(),
525                "injected".into(),
526            )]));
527            provider
528                .stream_simple(
529                    &requested,
530                    &Context::new(vec![]),
531                    &SimpleStreamOptions::default(),
532                )
533                .await
534                .result()
535                .await
536                .unwrap();
537            let recorded = calls.lock().unwrap().clone();
538            assert_eq!(recorded.len(), 1);
539            assert_eq!(recorded[0].0.base_url, "https://configured.example/v1");
540            assert_eq!(
541                recorded[0].0.headers.as_ref().unwrap()["x-api-key"],
542                "private-header"
543            );
544            assert_eq!(recorded[0].1.api_key.as_deref(), Some("private-credential"));
545            assert!(!recorded[0]
546                .0
547                .headers
548                .as_ref()
549                .unwrap()
550                .contains_key("authorization"));
551            drop(recorded);
552            requested.provider = "another-account".into();
553            let mut rejected = provider
554                .stream_simple(
555                    &requested,
556                    &Context::new(vec![]),
557                    &SimpleStreamOptions::default(),
558                )
559                .await;
560            assert!(matches!(
561                rejected.next().await,
562                Some(AssistantMessageEvent::Error {
563                    reason: ErrorReason::Error,
564                    ..
565                })
566            ));
567            assert_eq!(
568                rejected.result().await.unwrap().stop_reason,
569                StopReason::Error
570            );
571            assert_eq!(calls.lock().unwrap().len(), 1);
572        }
573    }
574
575    #[tokio::test]
576    async fn request_overrides_and_model_headers_remain_scoped_to_their_execution() {
577        let calls = Calls::default();
578        let mut factories = RuntimeFactories::default();
579        factories.providers.insert(
580            ("gateway".into(), "anthropic-messages".into()),
581            Arc::new(RecordingFactory(calls.clone())),
582        );
583        let mut models = vec![];
584        for (id, token) in [("first", "first-token"), ("second", "second-token")] {
585            let mut model = Model::new(
586                id,
587                id,
588                Api::AnthropicMessages,
589                "gateway",
590                "https://gateway.example",
591            );
592            model.headers = Some(BTreeMap::from([("x-api-key".into(), token.into())]));
593            models.push(model);
594        }
595        let (provider, _) = factories
596            .configure(
597                "gateway",
598                &Api::AnthropicMessages,
599                models,
600                RequestAuth::Headers(BTreeMap::from([(
601                    "authorization".into(),
602                    "Bearer default".into(),
603                )])),
604                &crate::RuntimeSettings::default(),
605                &BTreeMap::<String, String>::new(),
606            )
607            .unwrap();
608        let options = SimpleStreamOptions {
609            headers: Some(BTreeMap::from([(
610                "Authorization".into(),
611                "Bearer per-request".into(),
612            )])),
613            ..Default::default()
614        };
615        for model in provider.models() {
616            provider
617                .stream_simple(model, &Context::new(vec![]), &options)
618                .await
619                .result()
620                .await
621                .unwrap();
622        }
623        let recorded = calls.lock().unwrap().clone();
624        assert_eq!(
625            recorded[0].0.headers.as_ref().unwrap()["x-api-key"],
626            "first-token"
627        );
628        assert_eq!(
629            recorded[1].0.headers.as_ref().unwrap()["x-api-key"],
630            "second-token"
631        );
632        assert_eq!(recorded[0].1.headers.as_ref().unwrap().len(), 1);
633        assert_eq!(
634            recorded[0].1.headers.as_ref().unwrap()["Authorization"],
635            "Bearer per-request"
636        );
637    }
638}
rpi_runtime/provider/factories.rs

rpi_runtime/provider/
factories.rs
