mini_chat/
module.rs

1use std::sync::{Arc, Mutex, OnceLock};
2
3use async_trait::async_trait;
4use authz_resolver_sdk::AuthZResolverClient;
5use mini_chat_sdk::{MiniChatAuditPluginSpecV1, MiniChatModelPolicyPluginSpecV1};
6use modkit::api::OpenApiRegistry;
7use modkit::contracts::RunnableCapability;
8use modkit::{DatabaseCapability, Module, ModuleCtx, RestApiCapability};
9use modkit_db::outbox::{Outbox, OutboxHandle, Partitions};
10use oagw_sdk::ServiceGatewayClientV1;
11use sea_orm_migration::MigrationTrait;
12use tokio_util::sync::CancellationToken;
13use tracing::info;
14use types_registry_sdk::{RegisterResult, TypesRegistryClient};
15
16use crate::api::rest::routes;
17use crate::domain::service::{AppServices as GenericAppServices, Repositories};
18use crate::infra::outbox::{InfraOutboxEnqueuer, UsageEventHandler};
19
20pub(crate) type AppServices = GenericAppServices<
21    TurnRepository,
22    MessageRepository,
23    QuotaUsageRepository,
24    ReactionRepository,
25    ChatRepository,
26    ThreadSummaryRepository,
27>;
28use crate::infra::db::repo::attachment_repo::AttachmentRepository;
29use crate::infra::db::repo::chat_repo::ChatRepository;
30use crate::infra::db::repo::message_repo::MessageRepository;
31use crate::infra::db::repo::quota_usage_repo::QuotaUsageRepository;
32use crate::infra::db::repo::reaction_repo::ReactionRepository;
33use crate::infra::db::repo::thread_summary_repo::ThreadSummaryRepository;
34use crate::infra::db::repo::turn_repo::TurnRepository;
35use crate::infra::db::repo::vector_store_repo::VectorStoreRepository;
36use crate::infra::llm::provider_resolver::ProviderResolver;
37use crate::infra::model_policy::ModelPolicyGateway;
38
39/// Default URL prefix for all mini-chat REST routes.
40pub const DEFAULT_URL_PREFIX: &str = "/mini-chat";
41
42/// The mini-chat module: multi-tenant AI chat with SSE streaming.
43#[modkit::module(
44    name = "mini-chat",
45    deps = ["types-registry", "authz-resolver", "oagw"],
46    capabilities = [db, rest, stateful],
47)]
48pub struct MiniChatModule {
49    service: OnceLock<Arc<AppServices>>,
50    url_prefix: OnceLock<String>,
51    outbox_handle: Mutex<Option<OutboxHandle>>,
52}
53
54impl Default for MiniChatModule {
55    fn default() -> Self {
56        Self {
57            service: OnceLock::new(),
58            url_prefix: OnceLock::new(),
59            outbox_handle: Mutex::new(None),
60        }
61    }
62}
63
64#[async_trait]
65impl Module for MiniChatModule {
66    async fn init(&self, ctx: &ModuleCtx) -> anyhow::Result<()> {
67        info!("Initializing {} module", Self::MODULE_NAME);
68
69        let cfg: crate::config::MiniChatConfig = ctx.config_expanded()?;
70        cfg.streaming
71            .validate()
72            .map_err(|e| anyhow::anyhow!("streaming config: {e}"))?;
73        cfg.estimation_budgets
74            .validate()
75            .map_err(|e| anyhow::anyhow!("estimation_budgets config: {e}"))?;
76        cfg.quota
77            .validate()
78            .map_err(|e| anyhow::anyhow!("quota config: {e}"))?;
79        cfg.outbox
80            .validate()
81            .map_err(|e| anyhow::anyhow!("outbox config: {e}"))?;
82        cfg.context
83            .validate()
84            .map_err(|e| anyhow::anyhow!("context config: {e}"))?;
85        for (id, entry) in &cfg.providers {
86            entry
87                .validate(id)
88                .map_err(|e| anyhow::anyhow!("providers config: {e}"))?;
89        }
90
91        let vendor = cfg.vendor.trim().to_owned();
92        if vendor.is_empty() {
93            return Err(anyhow::anyhow!(
94                "{}: vendor must be a non-empty string",
95                Self::MODULE_NAME
96            ));
97        }
98
99        let registry = ctx.client_hub().get::<dyn TypesRegistryClient>()?;
100        register_plugin_schemas(
101            &*registry,
102            &[
103                (
104                    MiniChatModelPolicyPluginSpecV1::gts_schema_with_refs_as_string(),
105                    MiniChatModelPolicyPluginSpecV1::gts_schema_id(),
106                    "model-policy",
107                ),
108                (
109                    MiniChatAuditPluginSpecV1::gts_schema_with_refs_as_string(),
110                    MiniChatAuditPluginSpecV1::gts_schema_id(),
111                    "audit",
112                ),
113            ],
114        )
115        .await?;
116
117        self.url_prefix
118            .set(cfg.url_prefix)
119            .map_err(|_| anyhow::anyhow!("{} url_prefix already set", Self::MODULE_NAME))?;
120
121        let db_provider = ctx.db_required()?;
122        let db = Arc::new(db_provider);
123
124        // Create the model-policy gateway early for both outbox handler and services.
125        let model_policy_gw = Arc::new(ModelPolicyGateway::new(ctx.client_hub(), vendor));
126
127        // Start the outbox pipeline eagerly in init() (migrations ran in phase 2, DB is ready).
128        // The framework guarantees stop() is called on init failure, so the pipeline
129        // will be shut down cleanly if any later init step errors.
130        // The handler resolves the plugin lazily on first message delivery,
131        // avoiding a hard dependency on plugin availability during init().
132        let outbox_db = db.db();
133        let num_partitions = cfg.outbox.num_partitions;
134        let queue_name = cfg.outbox.queue_name.clone();
135
136        let outbox_handle =
137            Outbox::builder(outbox_db)
138                .queue(
139                    &queue_name,
140                    Partitions::of(u16::try_from(num_partitions).map_err(|_| {
141                        anyhow::anyhow!("num_partitions {num_partitions} exceeds u16")
142                    })?),
143                )
144                .decoupled(UsageEventHandler {
145                    plugin_provider: model_policy_gw.clone(),
146                })
147                .start()
148                .await
149                .map_err(|e| anyhow::anyhow!("outbox start: {e}"))?;
150
151        let outbox = Arc::clone(outbox_handle.outbox());
152        let outbox_enqueuer =
153            Arc::new(InfraOutboxEnqueuer::new(outbox, queue_name, num_partitions));
154
155        {
156            let mut guard = self
157                .outbox_handle
158                .lock()
159                .map_err(|e| anyhow::anyhow!("outbox_handle lock: {e}"))?;
160            if guard.is_some() {
161                return Err(anyhow::anyhow!(
162                    "{} outbox_handle already set",
163                    Self::MODULE_NAME
164                ));
165            }
166            *guard = Some(outbox_handle);
167        }
168
169        info!("Outbox pipeline started");
170
171        let authz = ctx
172            .client_hub()
173            .get::<dyn AuthZResolverClient>()
174            .map_err(|e| anyhow::anyhow!("failed to get AuthZ resolver: {e}"))?;
175
176        let gateway = ctx
177            .client_hub()
178            .get::<dyn ServiceGatewayClientV1>()
179            .map_err(|e| anyhow::anyhow!("failed to get OAGW gateway: {e}"))?;
180
181        // Register OAGW upstreams for each configured provider.
182        crate::infra::oagw_provisioning::register_oagw_upstreams(&gateway, &cfg.providers).await?;
183
184        let provider_resolver = Arc::new(ProviderResolver::new(&gateway, cfg.providers));
185
186        let repos = Repositories {
187            chat: Arc::new(ChatRepository::new(modkit_db::odata::LimitCfg {
188                default: 20,
189                max: 100,
190            })),
191            attachment: Arc::new(AttachmentRepository),
192            message: Arc::new(MessageRepository::new(modkit_db::odata::LimitCfg {
193                default: 20,
194                max: 100,
195            })),
196            quota: Arc::new(QuotaUsageRepository),
197            turn: Arc::new(TurnRepository),
198            reaction: Arc::new(ReactionRepository),
199            thread_summary: Arc::new(ThreadSummaryRepository),
200            vector_store: Arc::new(VectorStoreRepository),
201        };
202
203        let services = Arc::new(AppServices::new(
204            &repos,
205            db,
206            authz,
207            &(model_policy_gw.clone() as Arc<dyn crate::domain::repos::ModelResolver>),
208            provider_resolver,
209            cfg.streaming,
210            model_policy_gw.clone() as Arc<dyn crate::domain::repos::PolicySnapshotProvider>,
211            model_policy_gw as Arc<dyn crate::domain::repos::UserLimitsProvider>,
212            cfg.estimation_budgets,
213            cfg.quota,
214            outbox_enqueuer,
215            cfg.context,
216        ));
217
218        self.service
219            .set(services)
220            .map_err(|_| anyhow::anyhow!("{} module already initialized", Self::MODULE_NAME))?;
221
222        info!("{} module initialized successfully", Self::MODULE_NAME);
223        Ok(())
224    }
225}
226
227impl DatabaseCapability for MiniChatModule {
228    fn migrations(&self) -> Vec<Box<dyn MigrationTrait>> {
229        use sea_orm_migration::MigratorTrait;
230        info!("Providing mini-chat database migrations");
231        let mut m = crate::infra::db::migrations::Migrator::migrations();
232        m.extend(modkit_db::outbox::outbox_migrations());
233        m
234    }
235}
236
237impl RestApiCapability for MiniChatModule {
238    fn register_rest(
239        &self,
240        _ctx: &ModuleCtx,
241        router: axum::Router,
242        openapi: &dyn OpenApiRegistry,
243    ) -> anyhow::Result<axum::Router> {
244        let services = self
245            .service
246            .get()
247            .ok_or_else(|| anyhow::anyhow!("{} not initialized", Self::MODULE_NAME))?;
248
249        info!("Registering mini-chat REST routes");
250        let prefix = self
251            .url_prefix
252            .get()
253            .ok_or_else(|| anyhow::anyhow!("{} not initialized (url_prefix)", Self::MODULE_NAME))?;
254
255        let router = routes::register_routes(router, openapi, Arc::clone(services), prefix);
256        info!("Mini-chat REST routes registered successfully");
257        Ok(router)
258    }
259}
260
261#[async_trait]
262impl RunnableCapability for MiniChatModule {
263    async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
264        // Outbox pipeline already started in init().
265        Ok(())
266    }
267
268    async fn stop(&self, cancel: CancellationToken) -> anyhow::Result<()> {
269        let handle = self
270            .outbox_handle
271            .lock()
272            .map_err(|e| anyhow::anyhow!("outbox_handle lock: {e}"))?
273            .take();
274        if let Some(handle) = handle {
275            info!("Stopping outbox pipeline");
276            tokio::select! {
277                () = handle.stop() => {
278                    info!("Outbox pipeline stopped");
279                }
280                () = cancel.cancelled() => {
281                    info!("Outbox pipeline stop cancelled by framework deadline");
282                }
283            }
284        }
285        Ok(())
286    }
287}
288
289async fn register_plugin_schemas(
290    registry: &dyn TypesRegistryClient,
291    schemas: &[(String, &str, &str)],
292) -> anyhow::Result<()> {
293    let mut payload = Vec::with_capacity(schemas.len());
294    for (schema_str, schema_id, _label) in schemas {
295        let mut schema_json: serde_json::Value = serde_json::from_str(schema_str)?;
296        let obj = schema_json
297            .as_object_mut()
298            .ok_or_else(|| anyhow::anyhow!("schema {schema_id} is not a JSON object"))?;
299        obj.insert(
300            "additionalProperties".to_owned(),
301            serde_json::Value::Bool(false),
302        );
303        payload.push(schema_json);
304    }
305    let results = registry.register(payload).await?;
306    RegisterResult::ensure_all_ok(&results)?;
307    for (_schema_str, schema_id, label) in schemas {
308        info!(schema_id = %schema_id, "Registered {label} plugin schema in types-registry");
309    }
310    Ok(())
311}
mini_chat/module.rs

mini_chat/
module.rs
