Skip to main content

mini_chat/infra/plugins/static_model_policy/
gear.rs

1use std::sync::{Arc, OnceLock};
2
3use async_trait::async_trait;
4use mini_chat_sdk::{MiniChatModelPolicyPluginClientV1, MiniChatModelPolicyPluginSpecV1};
5use toolkit::Gear;
6use toolkit::client_hub::ClientScope;
7use toolkit::context::GearCtx;
8use toolkit::gts::PluginV1;
9use tracing::info;
10use types_registry_sdk::{RegisterResult, TypesRegistryClient};
11
12use super::config::StaticMiniChatPolicyPluginConfig;
13use super::service::Service;
14
15/// Static model-policy plugin gear for mini-chat.
16///
17/// Provides a config-driven model catalog for development and testing.
18#[toolkit::gear(
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 Gear for StaticMiniChatModelPolicyPlugin {
36    async fn init(&self, ctx: &GearCtx) -> anyhow::Result<()> {
37        let cfg: StaticMiniChatPolicyPluginConfig = ctx.config_or_default()?;
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!("{} gear already initialized", Self::MODULE_NAME))?;
56
57        // Build registration payload and instance id for this plugin.
58        let (instance_id, instance_json) =
59            PluginV1::<MiniChatModelPolicyPluginSpecV1>::build_registration(
60                "cf.core._.static_mini_chat_model_policy.v1",
61                cfg.vendor.clone(),
62                cfg.priority,
63            )?;
64
65        // Publish to types-registry.
66        let registry = ctx.client_hub().get::<dyn TypesRegistryClient>()?;
67        let results = registry.register(vec![instance_json]).await?;
68        RegisterResult::ensure_all_ok(&results)?;
69
70        // Register scoped client in ClientHub
71        let api: Arc<dyn MiniChatModelPolicyPluginClientV1> = service;
72        ctx.client_hub()
73            .register_scoped::<dyn MiniChatModelPolicyPluginClientV1>(
74                ClientScope::gts_id(&instance_id),
75                api,
76            );
77
78        info!(instance_id = %instance_id, "Static mini-chat model policy plugin registered");
79        Ok(())
80    }
81}