1use serde::de::DeserializeOwned;
2use std::sync::Arc;
3use tokio_util::sync::CancellationToken;
4use uuid::Uuid;
5
6use crate::{
8 config::{ConfigError, ConfigProvider, gear_config_or_default},
9 gear_config_required,
10};
11
12#[cfg(feature = "db")]
13pub(crate) type DbManager = toolkit_db::DbManager;
14#[cfg(feature = "db")]
15pub(crate) type DbProvider = toolkit_db::DBProvider<toolkit_db::DbError>;
16
17#[derive(Clone)]
18#[must_use]
19pub struct GearCtx {
20 gear_name: Arc<str>,
21 instance_id: Uuid,
22 config_provider: Arc<dyn ConfigProvider>,
23 client_hub: Arc<crate::client_hub::ClientHub>,
24 cancellation_token: CancellationToken,
25 #[cfg(feature = "db")]
26 db: Option<DbProvider>,
27}
28
29#[must_use]
34pub struct GearContextBuilder {
35 instance_id: Uuid,
36 config_provider: Arc<dyn ConfigProvider>,
37 client_hub: Arc<crate::client_hub::ClientHub>,
38 root_token: CancellationToken,
39 #[cfg(feature = "db")]
40 db_manager: Option<Arc<DbManager>>, }
42
43impl GearContextBuilder {
44 pub fn new(
45 instance_id: Uuid,
46 config_provider: Arc<dyn ConfigProvider>,
47 client_hub: Arc<crate::client_hub::ClientHub>,
48 root_token: CancellationToken,
49 ) -> Self {
50 Self {
51 instance_id,
52 config_provider,
53 client_hub,
54 root_token,
55 #[cfg(feature = "db")]
56 db_manager: None,
57 }
58 }
59
60 #[cfg(feature = "db")]
63 pub fn with_db_manager(mut self, db_manager: Arc<DbManager>) -> Self {
64 self.db_manager = Some(db_manager);
65 self
66 }
67
68 #[must_use]
70 pub fn instance_id(&self) -> Uuid {
71 self.instance_id
72 }
73
74 #[cfg_attr(not(feature = "db"), allow(clippy::unused_async))]
84 pub async fn for_gear(&self, gear_name: &str) -> anyhow::Result<GearCtx> {
85 let ctx = GearCtx::new(
86 Arc::<str>::from(gear_name),
87 self.instance_id,
88 self.config_provider.clone(),
89 self.client_hub.clone(),
90 self.root_token.child_token(),
91 );
92 #[cfg(feature = "db")]
93 let ctx = if let Some(mgr) = &self.db_manager
94 && let Some(handle) = mgr.get(gear_name).await?
95 {
96 ctx.with_db(toolkit_db::DBProvider::new(handle))
97 } else {
98 ctx
99 };
100 Ok(ctx)
101 }
102}
103
104impl GearCtx {
105 pub fn new(
109 gear_name: impl Into<Arc<str>>,
110 instance_id: Uuid,
111 config_provider: Arc<dyn ConfigProvider>,
112 client_hub: Arc<crate::client_hub::ClientHub>,
113 cancellation_token: CancellationToken,
114 ) -> Self {
115 Self {
116 gear_name: gear_name.into(),
117 instance_id,
118 config_provider,
119 client_hub,
120 cancellation_token,
121 #[cfg(feature = "db")]
122 db: None,
123 }
124 }
125
126 #[cfg(feature = "db")]
128 pub fn with_db(mut self, db: DbProvider) -> Self {
129 self.db = Some(db);
130 self
131 }
132
133 #[inline]
136 #[must_use]
137 pub fn gear_name(&self) -> &str {
138 &self.gear_name
139 }
140
141 #[inline]
146 #[must_use]
147 pub fn instance_id(&self) -> Uuid {
148 self.instance_id
149 }
150
151 #[inline]
152 #[must_use]
153 pub fn config_provider(&self) -> &dyn ConfigProvider {
154 &*self.config_provider
155 }
156
157 #[inline]
159 #[must_use]
160 pub fn client_hub(&self) -> Arc<crate::client_hub::ClientHub> {
161 self.client_hub.clone()
162 }
163
164 #[inline]
165 #[must_use]
166 pub fn cancellation_token(&self) -> &CancellationToken {
167 &self.cancellation_token
168 }
169
170 #[must_use]
189 #[cfg(feature = "db")]
190 pub fn db(&self) -> Option<toolkit_db::DBProvider<toolkit_db::DbError>> {
191 self.db.clone()
192 }
193
194 #[cfg(feature = "db")]
211 pub fn db_required(&self) -> anyhow::Result<toolkit_db::DBProvider<toolkit_db::DbError>> {
212 self.db().ok_or_else(|| {
213 anyhow::anyhow!("Database is not configured for gear '{}'", self.gear_name)
214 })
215 }
216
217 pub fn config<T: DeserializeOwned>(&self) -> Result<T, ConfigError> {
225 gear_config_required(self.config_provider.as_ref(), &self.gear_name)
226 }
227
228 pub fn config_or_default<T: DeserializeOwned + Default>(&self) -> Result<T, ConfigError> {
251 gear_config_or_default(self.config_provider.as_ref(), &self.gear_name)
252 }
253
254 pub fn config_expanded<T>(&self) -> Result<T, ConfigError>
261 where
262 T: DeserializeOwned + crate::var_expand::ExpandVars,
263 {
264 let mut cfg: T = self.config()?;
265 cfg.expand_vars().map_err(|e| ConfigError::VarExpand {
266 gear: self.gear_name.to_string(),
267 source: e,
268 })?;
269 Ok(cfg)
270 }
271
272 pub fn config_expanded_or_default<T>(&self) -> Result<T, ConfigError>
296 where
297 T: DeserializeOwned + Default + crate::var_expand::ExpandVars,
298 {
299 let mut cfg: T = self.config_or_default()?;
300 cfg.expand_vars().map_err(|e| ConfigError::VarExpand {
301 gear: self.gear_name.to_string(),
302 source: e,
303 })?;
304 Ok(cfg)
305 }
306
307 #[must_use]
310 pub fn raw_config(&self) -> &serde_json::Value {
311 use std::sync::LazyLock;
312
313 static EMPTY: LazyLock<serde_json::Value> =
314 LazyLock::new(|| serde_json::Value::Object(serde_json::Map::new()));
315
316 if let Some(gear_raw) = self.config_provider.get_gear_config(&self.gear_name) {
317 if let Some(obj) = gear_raw.as_object()
319 && let Some(config_section) = obj.get("config")
320 {
321 return config_section;
322 }
323 }
324 &EMPTY
325 }
326
327 pub fn without_db(&self) -> GearCtx {
330 GearCtx {
331 gear_name: self.gear_name.clone(),
332 instance_id: self.instance_id,
333 config_provider: self.config_provider.clone(),
334 client_hub: self.client_hub.clone(),
335 cancellation_token: self.cancellation_token.clone(),
336 #[cfg(feature = "db")]
337 db: None,
338 }
339 }
340}
341
342#[cfg(test)]
343#[cfg_attr(coverage_nightly, coverage(off))]
344mod tests {
345 use super::*;
346 use serde::Deserialize;
347 use serde_json::json;
348 use std::collections::HashMap;
349
350 #[derive(Debug, PartialEq, Deserialize, Default)]
351 struct TestConfig {
352 #[serde(default)]
353 api_key: String,
354 #[serde(default)]
355 timeout_ms: u64,
356 #[serde(default)]
357 enabled: bool,
358 }
359
360 struct MockConfigProvider {
361 gears: HashMap<String, serde_json::Value>,
362 }
363
364 impl MockConfigProvider {
365 fn new() -> Self {
366 let mut gears = HashMap::new();
367
368 gears.insert(
370 "test_gear".to_owned(),
371 json!({
372 "database": {
373 "url": "postgres://localhost/test"
374 },
375 "config": {
376 "api_key": "secret123",
377 "timeout_ms": 5000,
378 "enabled": true
379 }
380 }),
381 );
382
383 Self { gears }
384 }
385 }
386
387 impl ConfigProvider for MockConfigProvider {
388 fn get_gear_config(&self, gear_name: &str) -> Option<&serde_json::Value> {
389 self.gears.get(gear_name)
390 }
391 }
392
393 #[test]
394 fn test_gear_ctx_config_with_valid_config() {
395 let provider = Arc::new(MockConfigProvider::new());
396 let ctx = GearCtx::new(
397 "test_gear",
398 Uuid::new_v4(),
399 provider,
400 Arc::new(crate::client_hub::ClientHub::default()),
401 CancellationToken::new(),
402 );
403
404 let result: Result<TestConfig, ConfigError> = ctx.config();
405 assert!(result.is_ok());
406
407 let config = result.unwrap();
408 assert_eq!(config.api_key, "secret123");
409 assert_eq!(config.timeout_ms, 5000);
410 assert!(config.enabled);
411 }
412
413 #[test]
414 fn test_gear_ctx_config_returns_error_for_missing_gear() {
415 let provider = Arc::new(MockConfigProvider::new());
416 let ctx = GearCtx::new(
417 "nonexistent_gear",
418 Uuid::new_v4(),
419 provider,
420 Arc::new(crate::client_hub::ClientHub::default()),
421 CancellationToken::new(),
422 );
423
424 let result: Result<TestConfig, ConfigError> = ctx.config();
425 assert!(matches!(
426 result,
427 Err(ConfigError::GearNotFound { ref gear }) if gear == "nonexistent_gear"
428 ));
429 }
430
431 #[test]
432 fn test_gear_ctx_config_or_default_returns_default_for_missing_gear() {
433 let provider = Arc::new(MockConfigProvider::new());
434 let ctx = GearCtx::new(
435 "nonexistent_gear",
436 Uuid::new_v4(),
437 provider,
438 Arc::new(crate::client_hub::ClientHub::default()),
439 CancellationToken::new(),
440 );
441
442 let result: Result<TestConfig, ConfigError> = ctx.config_or_default();
443 assert!(result.is_ok());
444
445 let config = result.unwrap();
446 assert_eq!(config, TestConfig::default());
447 }
448
449 #[test]
450 fn test_gear_ctx_instance_id() {
451 let provider = Arc::new(MockConfigProvider::new());
452 let instance_id = Uuid::new_v4();
453 let ctx = GearCtx::new(
454 "test_gear",
455 instance_id,
456 provider,
457 Arc::new(crate::client_hub::ClientHub::default()),
458 CancellationToken::new(),
459 );
460
461 assert_eq!(ctx.instance_id(), instance_id);
462 }
463
464 #[derive(Debug, PartialEq, Deserialize, Default, toolkit_macros::ExpandVars)]
467 struct ExpandableConfig {
468 #[expand_vars]
469 #[serde(default)]
470 api_key: String,
471 #[expand_vars]
472 #[serde(default)]
473 endpoint: Option<String>,
474 #[serde(default)]
475 retries: u32,
476 }
477
478 fn make_ctx(gear_name: &str, config_json: serde_json::Value) -> GearCtx {
479 let mut gears = HashMap::new();
480 gears.insert(gear_name.to_owned(), config_json);
481 let provider = Arc::new(MockConfigProvider { gears });
482 GearCtx::new(
483 gear_name,
484 Uuid::new_v4(),
485 provider,
486 Arc::new(crate::client_hub::ClientHub::default()),
487 CancellationToken::new(),
488 )
489 }
490
491 #[test]
492 fn config_expanded_resolves_env_vars() {
493 let ctx = make_ctx(
494 "expand_mod",
495 json!({
496 "config": {
497 "api_key": "${TOOLKIT_TEST_KEY}",
498 "endpoint": "https://${TOOLKIT_TEST_HOST}/api",
499 "retries": 3
500 }
501 }),
502 );
503
504 temp_env::with_vars(
505 [
506 ("TOOLKIT_TEST_KEY", Some("secret-42")),
507 ("TOOLKIT_TEST_HOST", Some("example.com")),
508 ],
509 || {
510 let cfg: ExpandableConfig = ctx.config_expanded().unwrap();
511 assert_eq!(cfg.api_key, "secret-42");
512 assert_eq!(cfg.endpoint.as_deref(), Some("https://example.com/api"));
513 assert_eq!(cfg.retries, 3);
514 },
515 );
516 }
517
518 #[test]
519 fn config_expanded_returns_error_on_missing_var() {
520 let ctx = make_ctx(
521 "expand_mod",
522 json!({
523 "config": {
524 "api_key": "${TOOLKIT_TEST_MISSING_VAR_XYZ}"
525 }
526 }),
527 );
528
529 temp_env::with_vars([("TOOLKIT_TEST_MISSING_VAR_XYZ", None::<&str>)], || {
530 let err = ctx.config_expanded::<ExpandableConfig>().unwrap_err();
531 assert!(
532 matches!(err, ConfigError::VarExpand { ref gear, .. } if gear == "expand_mod"),
533 "expected EnvExpand error, got: {err:?}"
534 );
535 });
536 }
537
538 #[test]
539 fn config_expanded_skips_none_option_fields() {
540 let ctx = make_ctx(
541 "expand_mod",
542 json!({
543 "config": {
544 "api_key": "literal-key",
545 "retries": 5
546 }
547 }),
548 );
549
550 let cfg: ExpandableConfig = ctx.config_expanded().unwrap();
551 assert_eq!(cfg.api_key, "literal-key");
552 assert_eq!(cfg.endpoint, None);
553 assert_eq!(cfg.retries, 5);
554 }
555
556 #[test]
557 fn config_expanded_returns_error_when_missing() {
558 let ctx = make_ctx("missing_mod", json!({}));
559 let err = ctx.config_expanded::<ExpandableConfig>().unwrap_err();
560 assert!(matches!(
561 err,
562 ConfigError::MissingConfigSection { ref gear } if gear == "missing_mod"
563 ));
564 }
565
566 #[test]
567 fn config_expanded_or_default_falls_back_to_default_when_missing() {
568 let ctx = make_ctx("missing_mod", json!({}));
569 let cfg: ExpandableConfig = ctx.config_expanded_or_default().unwrap();
570 assert_eq!(cfg, ExpandableConfig::default());
571 }
572
573 #[derive(Debug, PartialEq, Deserialize, Default, toolkit_macros::ExpandVars)]
576 struct NestedProvider {
577 #[expand_vars]
578 #[serde(default)]
579 host: String,
580 #[expand_vars]
581 #[serde(default)]
582 token: Option<String>,
583 #[expand_vars]
584 #[serde(default)]
585 auth_config: Option<HashMap<String, String>>,
586 #[serde(default)]
587 port: u16,
588 }
589
590 #[derive(Debug, PartialEq, Deserialize, Default, toolkit_macros::ExpandVars)]
591 struct NestedConfig {
592 #[expand_vars]
593 #[serde(default)]
594 name: String,
595 #[expand_vars]
596 #[serde(default)]
597 providers: HashMap<String, NestedProvider>,
598 #[expand_vars]
599 #[serde(default)]
600 tags: Vec<String>,
601 }
602
603 #[test]
604 fn config_expanded_resolves_nested_structs() {
605 let ctx = make_ctx(
606 "nested_mod",
607 json!({
608 "config": {
609 "name": "${TOOLKIT_NESTED_NAME}",
610 "providers": {
611 "primary": {
612 "host": "${TOOLKIT_NESTED_HOST}",
613 "token": "${TOOLKIT_NESTED_TOKEN}",
614 "auth_config": {
615 "header": "X-Api-Key",
616 "secret_ref": "${TOOLKIT_NESTED_SECRET}"
617 },
618 "port": 443
619 }
620 },
621 "tags": ["${TOOLKIT_NESTED_TAG}", "literal"]
622 }
623 }),
624 );
625
626 temp_env::with_vars(
627 [
628 ("TOOLKIT_NESTED_NAME", Some("my-service")),
629 ("TOOLKIT_NESTED_HOST", Some("api.example.com")),
630 ("TOOLKIT_NESTED_TOKEN", Some("sk-secret")),
631 ("TOOLKIT_NESTED_SECRET", Some("key-12345")),
632 ("TOOLKIT_NESTED_TAG", Some("production")),
633 ],
634 || {
635 let cfg: NestedConfig = ctx.config_expanded().unwrap();
636 assert_eq!(cfg.name, "my-service");
637 assert_eq!(cfg.tags, vec!["production", "literal"]);
638
639 let primary = cfg.providers.get("primary").expect("primary provider");
640 assert_eq!(primary.host, "api.example.com");
641 assert_eq!(primary.token.as_deref(), Some("sk-secret"));
642 assert_eq!(primary.port, 443);
643
644 let auth = primary.auth_config.as_ref().expect("auth_config present");
645 assert_eq!(auth.get("header").map(String::as_str), Some("X-Api-Key"));
646 assert_eq!(
647 auth.get("secret_ref").map(String::as_str),
648 Some("key-12345")
649 );
650 },
651 );
652 }
653
654 #[test]
655 fn config_expanded_nested_missing_var_returns_error() {
656 let ctx = make_ctx(
657 "nested_mod",
658 json!({
659 "config": {
660 "name": "ok",
661 "providers": {
662 "bad": { "host": "${TOOLKIT_NESTED_GONE}", "port": 80 }
663 }
664 }
665 }),
666 );
667
668 temp_env::with_vars([("TOOLKIT_NESTED_GONE", None::<&str>)], || {
669 let err = ctx.config_expanded::<NestedConfig>().unwrap_err();
670 assert!(
671 matches!(err, ConfigError::VarExpand { ref gear, .. } if gear == "nested_mod"),
672 "expected EnvExpand, got: {err:?}"
673 );
674 });
675 }
676}