feat: add worker-backed PII local models

crates/cli/src/server/mod.rs

@@ -25,6 +25,7 @@ use nemo_relay::plugin::dynamic::{
 };
 use nemo_relay::plugin::{
     PluginComponentSpec, PluginConfig, clear_plugin_configuration, initialize_plugins_exact,
+    initialize_plugins_exact_with_inference_providers,
 };
 use nemo_relay_adaptive::plugin_component::register_adaptive_component;
 use nemo_relay_pii_redaction::component::register_pii_redaction_component;
@@ -1116,7 +1117,11 @@ impl PluginActivation {
                     CliError::Config(format!("worker plugin load failed: {error}"))
                 })?)
             };
-        initialize_plugins_exact(plugin_config)
+        let inference_providers = worker
+            .as_ref()
+            .map(WorkerPluginActivation::inference_providers)
+            .unwrap_or_default();
+        initialize_plugins_exact_with_inference_providers(plugin_config, inference_providers)
             .await
             .map_err(|error| CliError::Config(format!("plugin activation failed: {error}")))?;
         Ok(Self {

crates/core/src/lib.rs

@@ -69,6 +69,9 @@ mod registry;
 pub mod shared_runtime;
 pub mod stream;
 
+#[cfg(test)]
+#[path = "../tests/unit/inference_provider_tests.rs"]
+mod inference_provider_tests;
 #[cfg(test)]
 #[path = "../tests/unit/types_tests.rs"]
 mod types_tests;

crates/core/src/plugin.rs

@@ -48,6 +48,12 @@ pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel};
 
 pub mod dynamic;
 pub use dynamic::*;
+mod inference;
+#[doc(hidden)]
+pub use inference::{
+    InferenceProvider, InferenceProviderDescriptor, InferenceProviderFn,
+    InferenceProviderRegistration, InferenceProviderRegistry,
+};
 
 type PluginMap = HashMap<String, RegisteredPlugin>;
 
@@ -342,6 +348,7 @@ impl PluginRegistration {
 pub struct PluginRegistrationContext {
     registrations: Vec<PluginRegistration>,
     namespace: Option<String>,
+    inference_providers: InferenceProviderRegistry,
 }
 
 impl PluginRegistrationContext {
@@ -355,9 +362,33 @@ impl PluginRegistrationContext {
         Self {
             registrations: vec![],
             namespace: Some(namespace.into()),
+            inference_providers: InferenceProviderRegistry::default(),
         }
     }
 
+    /// Creates a registration context backed by host-owned inference providers.
+    #[doc(hidden)]
+    pub fn with_inference_providers(
+        namespace: Option<String>,
+        inference_providers: InferenceProviderRegistry,
+    ) -> Self {
+        Self {
+            registrations: Vec::new(),
+            namespace,
+            inference_providers,
+        }
+    }
+
+    /// Resolves an inference provider implementing the required contract.
+    #[doc(hidden)]
+    pub fn inference_provider(
+        &self,
+        name: &str,
+        expected_contract: &str,
+    ) -> Result<InferenceProvider> {
+        self.inference_providers.resolve(name, expected_contract)
+    }
+
     /// Returns the runtime-qualified name for a plugin-local registration.
     ///
     /// Plugin handlers should pass stable component-local names such as
@@ -1309,12 +1340,23 @@ pub fn plugin_config_schema() -> Json {
 /// is removed before the new configuration is activated.
 #[doc(hidden)]
 pub async fn initialize_plugins_exact(config: PluginConfig) -> Result<ConfigReport> {
+    initialize_plugins_exact_with_inference_providers(config, InferenceProviderRegistry::default())
+        .await
+}
+
+/// Configures plugin components with host-owned inference providers.
+#[doc(hidden)]
+pub async fn initialize_plugins_exact_with_inference_providers(
+    config: PluginConfig,
+    inference_providers: InferenceProviderRegistry,
+) -> Result<ConfigReport> {
     run_owned_plugin_mutation("plugin initialization", move || async move {
         let lease = LegacyPluginMutationLease::acquire()?;
         let rollback_failures = Arc::new(Mutex::new(Vec::new()));
         let initialization = tokio::spawn(initialize_plugins_exact_inner(
             config,
             Some(Arc::clone(&rollback_failures)),
+            inference_providers,
         ))
         .await
         .map_err(|error| {
@@ -1431,14 +1473,16 @@ pub(crate) async fn initialize_plugins_exact_for_host(
     config: PluginConfig,
     owner_id: u64,
     rollback_failures: Arc<Mutex<Vec<String>>>,
+    inference_providers: InferenceProviderRegistry,
 ) -> Result<ConfigReport> {
     verify_plugin_host_owner(owner_id)?;
-    initialize_plugins_exact_inner(config, Some(rollback_failures)).await
+    initialize_plugins_exact_inner(config, Some(rollback_failures), inference_providers).await
 }
 
 async fn initialize_plugins_exact_inner(
     config: PluginConfig,
     rollback_failures: Option<Arc<Mutex<Vec<String>>>>,
+    inference_providers: InferenceProviderRegistry,
 ) -> Result<ConfigReport> {
     let enabled_component_count = config
         .components
@@ -1475,11 +1519,17 @@ async fn initialize_plugins_exact_inner(
         match initialize_plugin_components_catching_panics(
             config.clone(),
             rollback_failures.clone(),
+            inference_providers.clone(),
         )
         .await
         {
             Ok(registrations) => {
-                store_active_plugin_configuration(config, report.clone(), registrations)?;
+                store_active_plugin_configuration(
+                    config,
+                    report.clone(),
+                    registrations,
+                    inference_providers,
+                )?;
                 log::info!(
                     target: "nemo_relay.plugin",
                     event = "plugin_configuration_replaced",
@@ -1491,6 +1541,7 @@ async fn initialize_plugins_exact_inner(
             Err(err) => match initialize_plugin_components_catching_panics(
                 previous_state.config.clone(),
                 rollback_failures.clone(),
+                previous_state.inference_providers.clone(),
             )
             .await
             {
@@ -1500,6 +1551,7 @@ async fn initialize_plugins_exact_inner(
                         previous_state.config,
                         previous_report,
                         registrations,
+                        previous_state.inference_providers,
                     )?;
                     log::warn!(
                         target: "nemo_relay.plugin",
@@ -1523,9 +1575,18 @@ async fn initialize_plugins_exact_inner(
             },
         }
     } else {
-        let registrations =
-            initialize_plugin_components_catching_panics(config.clone(), rollback_failures).await?;
-        store_active_plugin_configuration(config, report.clone(), registrations)?;
+        let registrations = initialize_plugin_components_catching_panics(
+            config.clone(),
+            rollback_failures,
+            inference_providers.clone(),
+        )
+        .await?;
+        store_active_plugin_configuration(
+            config,
+            report.clone(),
+            registrations,
+            inference_providers,
+        )?;
         log::info!(
             target: "nemo_relay.plugin",
             event = "plugin_configuration_activated",
@@ -1539,14 +1600,17 @@ async fn initialize_plugins_exact_inner(
 async fn initialize_plugin_components_catching_panics(
     config: PluginConfig,
     rollback_failures: Option<Arc<Mutex<Vec<String>>>>,
+    inference_providers: InferenceProviderRegistry,
 ) -> Result<Vec<PluginRegistration>> {
-    tokio::spawn(async move { initialize_plugin_components(&config, rollback_failures).await })
-        .await
-        .map_err(|error| {
-            PluginError::Internal(format!(
-                "plugin component initialization task failed: {error}"
-            ))
-        })?
+    tokio::spawn(async move {
+        initialize_plugin_components(&config, rollback_failures, inference_providers).await
+    })
+    .await
+    .map_err(|error| {
+        PluginError::Internal(format!(
+            "plugin component initialization task failed: {error}"
+        ))
+    })?
 }
 
 /// Validates and activates `config` layered on top of the discovered
@@ -1559,6 +1623,16 @@ pub async fn initialize_plugins(config: PluginConfig) -> Result<ConfigReport> {
     initialize_plugins_exact(config).await
 }
 
+/// Resolves discovered configuration and activates it with host inference providers.
+#[doc(hidden)]
+pub async fn initialize_plugins_with_inference_providers(
+    config: PluginConfig,
+    inference_providers: InferenceProviderRegistry,
+) -> Result<ConfigReport> {
+    let config = resolve_plugin_config(config)?;
+    initialize_plugins_exact_with_inference_providers(config, inference_providers).await
+}
+
 /// Layers `config` over the default discovered `plugins.toml` files.
 ///
 /// This is crate-visible so owned dynamic-plugin activation can use the same
@@ -1955,11 +2029,13 @@ struct ActivePluginConfiguration {
     config: PluginConfig,
     report: ConfigReport,
     registrations: Vec<PluginRegistration>,
+    inference_providers: InferenceProviderRegistry,
 }
 
 async fn initialize_plugin_components(
     config: &PluginConfig,
     rollback_failures: Option<Arc<Mutex<Vec<String>>>>,
+    inference_providers: InferenceProviderRegistry,
 ) -> Result<Vec<PluginRegistration>> {
     ensure_builtin_plugins_registered()?;
     let totals = plugin_component_totals(config);
@@ -1988,8 +2064,11 @@ async fn initialize_plugin_components(
             totals.get(component.kind.as_str()).copied().unwrap_or(1),
         );
 
-        let mut pending =
-            PendingPluginRegistrationContext::new(namespace, rollback_failures.clone());
+        let mut pending = PendingPluginRegistrationContext::new(
+            namespace,
+            rollback_failures.clone(),
+            inference_providers.clone(),
+        );
         plugin
             .register(&component.config, &mut pending.context)
             .await?;
@@ -2034,9 +2113,16 @@ struct PendingPluginRegistrationContext {
 }
 
 impl PendingPluginRegistrationContext {
-    fn new(namespace: String, rollback_failures: Option<Arc<Mutex<Vec<String>>>>) -> Self {
+    fn new(
+        namespace: String,
+        rollback_failures: Option<Arc<Mutex<Vec<String>>>>,
+        inference_providers: InferenceProviderRegistry,
+    ) -> Self {
         Self {
-            context: PluginRegistrationContext::with_namespace(namespace),
+            context: PluginRegistrationContext::with_inference_providers(
+                Some(namespace),
+                inference_providers,
+            ),
             rollback_failures,
         }
     }
@@ -2071,6 +2157,7 @@ fn store_active_plugin_configuration(
     config: PluginConfig,
     report: ConfigReport,
     registrations: Vec<PluginRegistration>,
+    inference_providers: InferenceProviderRegistry,
 ) -> Result<()> {
     let mut guard = ACTIVE_PLUGIN_CONFIGURATION.lock().map_err(|err| {
         PluginError::Internal(format!("active plugin configuration lock poisoned: {err}"))
@@ -2079,6 +2166,7 @@ fn store_active_plugin_configuration(
         config,
         report,
         registrations,
+        inference_providers,
     });
     Ok(())
 }

crates/core/src/plugin/dynamic/host.rs

@@ -16,8 +16,8 @@ use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value as Json};
 
 use crate::plugin::{
-    ConfigReport, PluginComponentSpec, PluginConfig, PluginHostLease, Result,
-    acquire_plugin_host_lease, clear_plugin_configuration_for_host,
+    ConfigReport, InferenceProviderRegistry, PluginComponentSpec, PluginConfig, PluginHostLease,
+    Result, acquire_plugin_host_lease, clear_plugin_configuration_for_host,
     ensure_builtin_plugins_registered, initialize_plugins_exact_for_host, resolve_plugin_config,
     run_owned_plugin_mutation,
 };
@@ -186,10 +186,18 @@ impl PluginHostActivation {
         );
         let rollback_failures = Arc::new(Mutex::new(Vec::new()));
         let owner_id = claim.owner_id();
+        #[cfg(feature = "worker-grpc")]
+        let inference_providers = worker
+            .as_ref()
+            .map(WorkerPluginActivation::inference_providers)
+            .unwrap_or_else(InferenceProviderRegistry::default);
+        #[cfg(not(feature = "worker-grpc"))]
+        let inference_providers = InferenceProviderRegistry::default();
         let initialization = tokio::spawn(initialize_plugins_exact_for_host(
             config,
             owner_id,
             Arc::clone(&rollback_failures),
+            inference_providers,
         ))
         .await
         .map_err(|error| {
@@ -262,6 +270,16 @@ impl PluginHostActivation {
         self.active
     }
 
+    /// Returns the inference providers owned by this activation.
+    #[doc(hidden)]
+    pub fn inference_providers(&self) -> InferenceProviderRegistry {
+        #[cfg(feature = "worker-grpc")]
+        if let Some(worker) = &self.worker {
+            return worker.inference_providers();
+        }
+        InferenceProviderRegistry::default()
+    }
+
     /// Clear registered callbacks before unloading libraries and workers.
     pub fn clear(mut self) -> Result<()> {
         self.clear_inner()
@@ -300,6 +318,7 @@ impl PluginHostActivation {
         #[cfg(feature = "worker-grpc")]
         if let Some(worker) = &mut self.worker {
             runtime_outcome.merge(worker.deregister_plugin_kinds_checked());
+            runtime_outcome.merge(worker.deregister_inference_providers_checked());
         }
 
         // A worker cannot be stopped while its registry adapter might still be