cc_switch/daemon/aggregate/
state.rs

1use crate::config::types::ConfigStorage;
2use crate::daemon::aggregate::stream::TaggedCaptureEvent;
3use ccs_proxy::store::FsStore;
4use std::collections::BTreeMap;
5use std::sync::Arc;
6use tokio::sync::broadcast;
7
8pub type StoreEntry = (String, Arc<FsStore>);
9pub type AliasEntry = (String, Vec<String>);
10
11pub struct AliasMap {
12    map: BTreeMap<String, Vec<String>>,
13}
14
15impl AliasMap {
16    pub fn from_storage(storage: &ConfigStorage) -> Self {
17        let mut map: BTreeMap<String, Vec<String>> = BTreeMap::new();
18        for config in storage.configurations.values() {
19            if !config.url.is_empty() {
20                map.entry(config.url.clone())
21                    .or_default()
22                    .push(config.alias_name.clone());
23            }
24        }
25        Self { map }
26    }
27
28    pub fn from_entries(entries: Vec<AliasEntry>) -> Self {
29        Self {
30            map: entries.into_iter().collect(),
31        }
32    }
33
34    pub fn aliases_for(&self, upstream: &str) -> Vec<String> {
35        self.map.get(upstream).cloned().unwrap_or_default()
36    }
37}
38
39pub struct AggregateState {
40    pub stores: Vec<StoreEntry>,
41    pub merged_events: broadcast::Sender<TaggedCaptureEvent>,
42    pub alias_map: Arc<AliasMap>,
43    pub started_at: chrono::DateTime<chrono::Utc>,
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49    use crate::config::types::{ConfigStorage, Configuration};
50    use std::collections::BTreeMap;
51
52    fn make_storage(entries: &[(&str, &str)]) -> ConfigStorage {
53        let mut configurations = BTreeMap::new();
54        for (alias, url) in entries {
55            configurations.insert(
56                alias.to_string(),
57                Configuration {
58                    alias_name: alias.to_string(),
59                    token: "sk-test".to_string(),
60                    url: url.to_string(),
61                    ..Default::default()
62                },
63            );
64        }
65        ConfigStorage {
66            configurations,
67            claude_settings_dir: None,
68            default_storage_mode: None,
69            codex_configurations: None,
70        }
71    }
72
73    #[test]
74    fn alias_map_groups_by_upstream() {
75        let storage = make_storage(&[
76            ("work", "https://api.anthropic.com"),
77            ("personal", "https://api.anthropic.com"),
78            ("other", "https://other.example.com"),
79        ]);
80        let map = AliasMap::from_storage(&storage);
81        let mut aliases = map.aliases_for("https://api.anthropic.com");
82        aliases.sort();
83        assert_eq!(aliases, vec!["personal", "work"]);
84        assert_eq!(map.aliases_for("https://other.example.com"), vec!["other"]);
85    }
86
87    #[test]
88    fn alias_map_returns_empty_for_unknown() {
89        let storage = make_storage(&[("work", "https://api.anthropic.com")]);
90        let map = AliasMap::from_storage(&storage);
91        assert!(map.aliases_for("https://unknown.example.com").is_empty());
92    }
93
94    #[test]
95    fn alias_map_from_entries() {
96        let map = AliasMap::from_entries(vec![(
97            "https://a.example.com".to_string(),
98            vec!["alias_a".to_string()],
99        )]);
100        assert_eq!(map.aliases_for("https://a.example.com"), vec!["alias_a"]);
101        assert!(map.aliases_for("https://unknown.com").is_empty());
102    }
103}
cc_switch/daemon/aggregate/state.rs

cc_switch/daemon/aggregate/
state.rs
