mj_controller/controller/
profile_config.rs

1//! Automatically populated, persistent profile capabilities.
2
3use mj_core::hex::lower_hex;
4use std::collections::BTreeMap;
5use std::sync::{Arc, Mutex, OnceLock, Weak};
6use std::time::Duration;
7
8use crate::targets::{CancellableProcessExecutor, CommandExecutor, CommandSpec, TargetLocator};
9use anyhow::{Context, Result, ensure};
10use mj_core::config::{Config, HarnessProfile};
11use mj_core::worker_launch::{ProfileConfig, ProfileProbeSpec};
12use sha2::{Digest, Sha256};
13
14#[derive(Default)]
15struct ProbeLock {
16    gate: tokio::sync::Mutex<()>,
17    cancelled: Arc<std::sync::atomic::AtomicBool>,
18}
19
20pub fn cancel_all() {
21    if let Some(probes) = PROBES.get() {
22        for probe in probes
23            .lock()
24            .expect("profile probe locks poisoned")
25            .values()
26            .filter_map(Weak::upgrade)
27        {
28            probe
29                .cancelled
30                .store(true, std::sync::atomic::Ordering::Release);
31        }
32    }
33}
34static PROBES: OnceLock<Mutex<BTreeMap<String, Weak<ProbeLock>>>> = OnceLock::new();
35
36/// A request owns its background job independently of the HTTP connection.
37/// Concurrent callers serialize by profile and reuse the completed cache.
38pub async fn discover(
39    profile_id: String,
40    model: Option<String>,
41    refresh: bool,
42) -> Result<ProfileConfig> {
43    serialized(profile_id.clone(), move |cancelled| {
44        discover_blocking(&profile_id, model, refresh, cancelled)
45    })
46    .await
47}
48
49async fn serialized(
50    profile_id: String,
51    job: impl FnOnce(Arc<std::sync::atomic::AtomicBool>) -> Result<ProfileConfig> + Send + 'static,
52) -> Result<ProfileConfig> {
53    let lock = {
54        let mut locks = PROBES
55            .get_or_init(Default::default)
56            .lock()
57            .expect("profile probe locks poisoned");
58        locks.retain(|_, lock| lock.strong_count() > 0);
59        let lock = locks
60            .get(&profile_id)
61            .and_then(Weak::upgrade)
62            .unwrap_or_else(|| Arc::new(ProbeLock::default()));
63        locks.insert(profile_id, Arc::downgrade(&lock));
64        lock
65    };
66    tokio::spawn(async move {
67        let _guard = lock.gate.lock().await;
68        let cancelled = lock.cancelled.clone();
69        let result = tokio::task::spawn_blocking(move || job(cancelled))
70            .await
71            .context("profile discovery task panicked")?;
72        if let Err(error) = &result {
73            tracing::warn!(error = %format!("{error:#}"), "profile discovery failed");
74        }
75        result
76    })
77    .await
78    .context("profile discovery supervisor panicked")?
79}
80
81/// Update a model-specific entry only when a managed worker matches the local
82/// probe binary. Container/ambient installations cannot establish that match.
83pub async fn observe(
84    profile_id: String,
85    worker_build: String,
86    state: mj_core::relay::RelayOperationalState,
87) -> Result<()> {
88    serialized(profile_id.clone(), move |cancelled| {
89        let config = Config::load()?;
90        let profile = config
91            .enabled_profile(&profile_id)
92            .context("observed profile was removed")?;
93        let facts = mj_core::acp::AcpSessionFacts::from_operational(
94            profile.kind,
95            &state.config,
96            &state.config_options,
97            state.modes.as_ref(),
98        );
99        let mut choices = ProfileConfig {
100            model: facts.current_model().map(str::to_owned),
101            models: mj_core::acp::session_config_choices(&state.config_options, "model"),
102            efforts: mj_core::acp::session_config_choices(&state.config_options, "effort"),
103            observed_at: chrono::Utc::now().timestamp(),
104        };
105        enrich_profile_config(profile, &mut choices)?;
106        // Worker observations do not carry authentication provenance. Claude's
107        // setup-token and login catalogues differ, so only probes may cache it.
108        if profile.kind == mj_core::config::HarnessKind::Claude {
109            return Ok(choices);
110        }
111        let executor =
112            CancellableProcessExecutor::new(cancelled).with_deadline(Duration::from_secs(30));
113        let worker = super::worker_binary::worker_binary_for(
114            &TargetLocator::LocalBare {
115                worker_root: String::new(),
116            },
117            &executor,
118        )?;
119        if mj_core::worker_launch::worker_executable_digest(&worker)? == worker_build {
120            store(
121                &profile_id,
122                &fingerprint(profile, &profile.environment)?,
123                &choices.model,
124                &choices,
125            )?;
126        }
127        Ok(choices)
128    })
129    .await
130    .map(|_| ())
131}
132
133fn fingerprint(profile: &HarnessProfile, environment: &BTreeMap<String, String>) -> Result<String> {
134    let mut hash = Sha256::new();
135    hash.update(b"profile-config-v3\0");
136    hash.update(serde_json::to_vec(profile)?);
137    hash.update(serde_json::to_vec(environment)?);
138    hash.update(
139        mj_core::harness_runtime::pin(profile.kind)
140            .install_id
141            .as_bytes(),
142    );
143    Ok(lower_hex(hash.finalize()))
144}
145
146fn discover_blocking(
147    profile_id: &str,
148    model: Option<String>,
149    refresh: bool,
150    cancelled: Arc<std::sync::atomic::AtomicBool>,
151) -> Result<ProfileConfig> {
152    ensure!(
153        !cancelled.load(std::sync::atomic::Ordering::Acquire),
154        "profile discovery cancelled"
155    );
156    let config = Config::load()?;
157    let profile = config
158        .enabled_profile(profile_id)
159        .with_context(|| format!("unknown or disabled profile {profile_id:?}"))?;
160    let mut environment = profile.environment.clone();
161    super::worker_binary::apply_claude_setup_token(
162        &mut environment,
163        profile.kind,
164        &mj_core::credentials::claude_oauth_token_path(profile_id),
165    );
166    let fingerprint = fingerprint(profile, &environment)?;
167    resolve_cached(
168        refresh,
169        || {
170            crate::database::load_profile_config_cache(
171                profile_id,
172                model.as_deref().unwrap_or_default(),
173                &fingerprint,
174            )?
175            .map(|body| serde_json::from_str(&body).context("read cached profile configuration"))
176            .transpose()
177        },
178        || {
179            let mut choices = probe_profile(
180                profile_id,
181                profile,
182                environment.clone(),
183                model.clone(),
184                cancelled,
185            )?;
186            enrich_profile_config(profile, &mut choices)?;
187            Ok(choices)
188        },
189        |choices| {
190            if model.is_none() || choices.model == model {
191                store(profile_id, &fingerprint, &model, choices)?;
192            }
193            store(profile_id, &fingerprint, &choices.model, choices)
194        },
195    )
196}
197
198/// Muse's ACP bridge reports effort choices but no model selector. Its native
199/// settings file is authoritative for the one model this profile will run, so
200/// publish that value rather than making callers create a session to learn it.
201fn enrich_profile_config(profile: &HarnessProfile, choices: &mut ProfileConfig) -> Result<()> {
202    if profile.kind != mj_core::config::HarnessKind::Muse || !choices.models.is_empty() {
203        return Ok(());
204    }
205    let path = profile.home.join("settings.json");
206    let metadata = std::fs::metadata(&path)
207        .with_context(|| format!("read Muse settings metadata {}", path.display()))?;
208    ensure!(metadata.len() <= 1024 * 1024, "Muse settings are too large");
209    let settings: serde_json::Value = serde_json::from_slice(
210        &std::fs::read(&path).with_context(|| format!("read Muse settings {}", path.display()))?,
211    )
212    .context("decode Muse settings")?;
213    let model = settings
214        .get("model")
215        .and_then(serde_json::Value::as_str)
216        .map(str::trim)
217        .filter(|model| !model.is_empty())
218        .context("Muse settings do not select a model")?;
219    choices.model = Some(model.to_owned());
220    choices.models.push(mj_core::acp::SessionConfigChoice {
221        value: model.to_owned(),
222        name: model.to_owned(),
223        description: Some("Configured by Muse Code settings".into()),
224    });
225    Ok(())
226}
227
228fn resolve_cached(
229    refresh: bool,
230    load: impl FnOnce() -> Result<Option<ProfileConfig>>,
231    probe: impl FnOnce() -> Result<ProfileConfig>,
232    save: impl FnOnce(&ProfileConfig) -> Result<()>,
233) -> Result<ProfileConfig> {
234    if !refresh && let Some(choices) = load()? {
235        return Ok(choices);
236    }
237    let choices = probe()?;
238    save(&choices)?;
239    Ok(choices)
240}
241
242fn probe_profile(
243    profile_id: &str,
244    profile: &HarnessProfile,
245    environment: BTreeMap<String, String>,
246    model: Option<String>,
247    cancelled: Arc<std::sync::atomic::AtomicBool>,
248) -> Result<ProfileConfig> {
249    let root = tempfile::tempdir().context("create private profile discovery directory")?;
250    let home = root.path().join("profile");
251    super::worker_binary::stage_profile(profile, &home)?;
252    super::worker_binary::stage_codex_catalog(
253        profile_id,
254        profile,
255        &home,
256        &super::worker_binary::fetch_catalog_over_https,
257        &super::worker_binary::SharedCatalogCache,
258    )?;
259    let cwd = root.path().join("workspace");
260    std::fs::create_dir(&cwd)?;
261    let executor =
262        CancellableProcessExecutor::new(cancelled).with_deadline(Duration::from_secs(300));
263    let worker = super::worker_binary::worker_binary_for(
264        &TargetLocator::LocalBare {
265            worker_root: root.path().to_string_lossy().into_owned(),
266        },
267        &executor,
268    )?;
269    let spec = ProfileProbeSpec {
270        harness: profile.kind,
271        profile_home: home,
272        environment,
273        cwd,
274        model,
275    };
276    let path = root.path().join("probe.json");
277    mj_core::config::atomic_write(&path, &serde_json::to_vec(&spec)?)?;
278    let command = CommandSpec::new(
279        worker.to_string_lossy(),
280        [
281            "worker".to_owned(),
282            "discover-config".into(),
283            "--spec".into(),
284            path.to_string_lossy().into_owned(),
285        ],
286    )
287    .purpose("discover profile configuration");
288    let output = executor.execute(&command)?;
289    if !output.stderr.is_empty() {
290        tracing::info!(harness = ?profile.kind, diagnostics = %String::from_utf8_lossy(&output.stderr).trim(), "profile discovery worker diagnostics");
291    }
292    ensure!(
293        output.status == 0,
294        "profile discovery failed: {}",
295        String::from_utf8_lossy(&output.stderr).trim()
296    );
297    let choices: ProfileConfig =
298        serde_json::from_slice(&output.stdout).context("decode discovered configuration")?;
299    Ok(choices)
300}
301
302fn store(
303    profile: &str,
304    fingerprint: &str,
305    model: &Option<String>,
306    choices: &ProfileConfig,
307) -> Result<()> {
308    crate::database::save_profile_config_cache(
309        profile.to_owned(),
310        model.clone().unwrap_or_default(),
311        fingerprint.to_owned(),
312        serde_json::to_string(choices)?,
313    )
314}
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319    use std::cell::{Cell, RefCell};
320
321    #[test]
322    fn muse_discovery_publishes_the_model_selected_by_native_settings() {
323        let home = tempfile::tempdir().unwrap();
324        std::fs::write(
325            home.path().join("settings.json"),
326            br#"{"model":"muse-spark-1.3-contributor"}"#,
327        )
328        .unwrap();
329        let profile = HarnessProfile {
330            enabled: true,
331            kind: mj_core::config::HarnessKind::Muse,
332            home: home.path().into(),
333            environment: BTreeMap::new(),
334            context_window_bytes: None,
335            guardian_review_model: None,
336        };
337        let mut choices = ProfileConfig {
338            model: Some(String::new()),
339            models: Vec::new(),
340            efforts: Vec::new(),
341            observed_at: 1,
342        };
343
344        enrich_profile_config(&profile, &mut choices).unwrap();
345
346        assert_eq!(choices.model.as_deref(), Some("muse-spark-1.3-contributor"));
347        assert_eq!(choices.models.len(), 1);
348        assert_eq!(choices.models[0].value, "muse-spark-1.3-contributor");
349    }
350
351    #[test]
352    fn setup_token_changes_invalidate_the_discovery_cache() {
353        use mj_core::credentials::{CLAUDE_OAUTH_TOKEN_ENV, write_claude_oauth_token};
354        let root = tempfile::tempdir().unwrap();
355        let token = root.path().join("token");
356        let profile = HarnessProfile {
357            enabled: true,
358            kind: mj_core::config::HarnessKind::Claude,
359            home: root.path().into(),
360            environment: BTreeMap::new(),
361            context_window_bytes: None,
362            guardian_review_model: None,
363        };
364        let resolve = |environment: BTreeMap<String, String>| {
365            let mut environment = environment;
366            super::super::worker_binary::apply_claude_setup_token(
367                &mut environment,
368                profile.kind,
369                &token,
370            );
371            environment
372        };
373        let login = fingerprint(&profile, &resolve(BTreeMap::new())).unwrap();
374        write_claude_oauth_token(&token, b"setup-first").unwrap();
375        let first = resolve(BTreeMap::new());
376        assert_eq!(first[CLAUDE_OAUTH_TOKEN_ENV], "setup-first");
377        let first_key = fingerprint(&profile, &first).unwrap();
378        assert_ne!(login, first_key);
379        write_claude_oauth_token(&token, b"setup-second").unwrap();
380        assert_eq!(
381            first[CLAUDE_OAUTH_TOKEN_ENV], "setup-first",
382            "an in-flight probe retains its authentication snapshot"
383        );
384        assert_ne!(
385            first_key,
386            fingerprint(&profile, &resolve(BTreeMap::new())).unwrap()
387        );
388        let explicit = BTreeMap::from([(CLAUDE_OAUTH_TOKEN_ENV.into(), "explicit".into())]);
389        assert_eq!(resolve(explicit.clone()), explicit);
390        assert!(!first_key.contains("setup-first"));
391    }
392
393    #[tokio::test]
394    async fn concurrent_cold_lookups_share_the_first_probe() {
395        let cache = Arc::new(Mutex::new(None));
396        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
397        let mut tasks = vec![];
398        for _ in 0..8 {
399            let cache = cache.clone();
400            let calls = calls.clone();
401            tasks.push(tokio::spawn(serialized(
402                "concurrent-cache-test".into(),
403                move |_| {
404                    resolve_cached(
405                        false,
406                        || Ok(cache.lock().unwrap().clone()),
407                        || {
408                            calls.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
409                            std::thread::sleep(Duration::from_millis(20));
410                            Ok(ProfileConfig {
411                                model: None,
412                                models: vec![],
413                                efforts: vec![],
414                                observed_at: 1,
415                            })
416                        },
417                        |value| {
418                            *cache.lock().unwrap() = Some(value.clone());
419                            Ok(())
420                        },
421                    )
422                },
423            )));
424        }
425        for task in tasks {
426            task.await.unwrap().unwrap();
427        }
428        assert_eq!(calls.load(std::sync::atomic::Ordering::Relaxed), 1);
429    }
430
431    #[test]
432    fn empty_cache_discovers_automatically_and_failed_probes_remain_retryable() {
433        let cache = RefCell::new(None);
434        let calls = Cell::new(0);
435        let choices = ProfileConfig {
436            model: Some("full/model".into()),
437            models: vec![],
438            efforts: vec![],
439            observed_at: 1,
440        };
441        let lookup = |fail: bool| {
442            resolve_cached(
443                false,
444                || Ok(cache.borrow().clone()),
445                || {
446                    calls.set(calls.get() + 1);
447                    ensure!(!fail, "probe failed");
448                    Ok(choices.clone())
449                },
450                |value| {
451                    *cache.borrow_mut() = Some(value.clone());
452                    Ok(())
453                },
454            )
455        };
456        assert!(lookup(true).is_err());
457        assert!(cache.borrow().is_none());
458        assert_eq!(lookup(false).unwrap(), choices);
459        assert_eq!(
460            lookup(true).unwrap(),
461            choices,
462            "a warm cache must not run the failing probe"
463        );
464        assert_eq!(calls.get(), 2);
465        assert_eq!(
466            resolve_cached(
467                true,
468                || Ok(cache.borrow().clone()),
469                || Ok(ProfileConfig {
470                    observed_at: 2,
471                    ..choices.clone()
472                }),
473                |_| Ok(())
474            )
475            .unwrap()
476            .observed_at,
477            2
478        );
479    }
480}
mj_controller/controller/profile_config.rs

mj_controller/controller/
profile_config.rs
