Skip to main content

mini_chat/infra/plugins/static_model_policy/
module.rs

1use std::sync::{Arc, OnceLock};
2
3use async_trait::async_trait;
4use mini_chat_sdk::{MiniChatModelPolicyPluginClientV1, MiniChatModelPolicyPluginSpecV1};
5use modkit::Module;
6use modkit::client_hub::ClientScope;
7use modkit::context::ModuleCtx;
8use modkit::gts::BaseModkitPluginV1;
9use tracing::info;
10use types_registry_sdk::{RegisterResult, TypesRegistryClient};
11
12use super::config::StaticMiniChatPolicyPluginConfig;
13use super::service::Service;
14
15/// Static model-policy plugin module for mini-chat.
16///
17/// Provides a config-driven model catalog for development and testing.
18#[modkit::module(
19    name = "static-mini-chat-model-policy-plugin",
20    deps = ["types-registry"]
21)]
22pub struct StaticMiniChatModelPolicyPlugin {
23    service: OnceLock<Arc<Service>>,
24}
25
26impl Default for StaticMiniChatModelPolicyPlugin {
27    fn default() -> Self {
28        Self {
29            service: OnceLock::new(),
30        }
31    }
32}
33
34#[async_trait]
35impl Module for StaticMiniChatModelPolicyPlugin {
36    async fn init(&self, ctx: &ModuleCtx) -> anyhow::Result<()> {
37        let cfg: StaticMiniChatPolicyPluginConfig = ctx.config()?;
38        info!(
39            vendor = %cfg.vendor,
40            priority = cfg.priority,
41            models = cfg.model_catalog.len(),
42            "Loaded static mini-chat model policy plugin configuration"
43        );
44
45        // Create service and lock initialization before any external side-effects
46        // so that retries fail fast without duplicate registrations.
47        let service = Arc::new(Service::new(
48            cfg.model_catalog,
49            cfg.kill_switches,
50            cfg.default_standard_limits,
51            cfg.default_premium_limits,
52        ));
53        self.service
54            .set(service.clone())
55            .map_err(|_| anyhow::anyhow!("{} module already initialized", Self::MODULE_NAME))?;
56
57        // Generate plugin instance ID
58        let instance_id = MiniChatModelPolicyPluginSpecV1::gts_make_instance_id(
59            "x.core._.static_mini_chat_model_policy.v1",
60        );
61
62        // Register plugin instance in types-registry
63        let registry = ctx.client_hub().get::<dyn TypesRegistryClient>()?;
64        let instance = BaseModkitPluginV1::<MiniChatModelPolicyPluginSpecV1> {
65            id: instance_id.clone(),
66            vendor: cfg.vendor.clone(),
67            priority: cfg.priority,
68            properties: MiniChatModelPolicyPluginSpecV1,
69        };
70        let instance_json = serde_json::to_value(&instance)?;
71
72        let results = registry.register(vec![instance_json]).await?;
73        RegisterResult::ensure_all_ok(&results)?;
74
75        // Register scoped client in ClientHub
76        let api: Arc<dyn MiniChatModelPolicyPluginClientV1> = service;
77        ctx.client_hub()
78            .register_scoped::<dyn MiniChatModelPolicyPluginClientV1>(
79                ClientScope::gts_id(&instance_id),
80                api,
81            );
82
83        info!(instance_id = %instance_id, "Static mini-chat model policy plugin registered");
84        Ok(())
85    }
86}