Skip to main content

agentmux_mcp/
self_keys.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Identity M4d-3: this agent's signing keys, fetched from srv
5//! (`GET /agentmux/agents/self/keys`) rather than only read from `.mcp.json`
6//! (`SPEC_AGENT_IDENTITY_CARRIED_NOT_DERIVED_2026_09_23.md` ยง6.5.10).
7//!
8//! The env copy goes stale: srv rotates the host jekt key after 24 h, and an
9//! agent running longer keeps signing with the key it was launched with, so
10//! every jekt it sends fails srv's check. Fetched at startup and again once
11//! the fetched key reaches the expiry srv served with it: srv rotates a key
12//! only after that, so until then nothing newer exists, and the fetch that
13//! follows it is what rotates it. Checked before every tool call, so a UI
14//! automation proof is signed with a live key too.
15//!
16//! The fetched name-keyed keys are used only when the slug srv served them
17//! for is exactly this process's `AGENTMUX_AGENT_ID`: an agent signing under
18//! a name its row doesn't have (a collision-suffixed backfill, a #3573 stub)
19//! keeps its env keys, which are the ones filed under the name it signs as.
20//! The UID-keyed keys are held for M4d-6's v2 signatures. Nothing here is
21//! logged.
22
23use std::sync::Mutex;
24use std::time::{Duration, Instant};
25
26use agentmux_common::AUTH_KEY_HEADER;
27use serde::Deserialize;
28
29/// Re-fetch a copy this old even with no expiry to go by (an older srv).
30const REFRESH_AFTER: Duration = Duration::from_secs(20 * 60 * 60);
31
32#[derive(Deserialize, Default, Clone)]
33struct Wire {
34    slug: String,
35    jekt_key: Option<String>,
36    /// Unix seconds; absent from an srv that predates it.
37    #[serde(default)]
38    jekt_key_expires_at: Option<i64>,
39    lan_key: Option<String>,
40    wan_key: Option<String>,
41    #[serde(default)]
42    uid: String,
43    uid_lan_key: Option<String>,
44}
45
46struct Cached {
47    keys: Wire,
48    fetched_at: Instant,
49}
50
51static CACHE: Mutex<Option<Cached>> = Mutex::new(None);
52
53/// When a fetch was last tried. With no usable copy (an srv without the
54/// endpoint, an unreachable one), every tool call would otherwise retry and
55/// wait out the request; this spaces the retries.
56static LAST_ATTEMPT: Mutex<Option<Instant>> = Mutex::new(None);
57const RETRY_AFTER: Duration = Duration::from_secs(60);
58/// Short, so a slow srv delays a tool call by this at most.
59const FETCH_TIMEOUT: Duration = Duration::from_secs(3);
60
61/// Fetch the keys if there is no copy yet, or it is older than 20 h. Errors
62/// leave the old copy (or none) in place: signing then falls back to the env
63/// keys, exactly as before M4d-3.
64pub(crate) async fn refresh_if_stale(client: &reqwest::Client, local_url: &str, auth_key: &str) {
65    let fresh = CACHE.lock().unwrap_or_else(|e| e.into_inner()).as_ref().is_some_and(|c| still_good(c, agentmux_common::time::now_secs()));
66    if fresh || local_url.is_empty() || std::env::var("AGENTMUX_AGENT_TOKEN").map_or(true, |t| t.trim().is_empty()) {
67        return;
68    }
69    {
70        let mut last = LAST_ATTEMPT.lock().unwrap_or_else(|e| e.into_inner());
71        if last.is_some_and(|t| t.elapsed() < RETRY_AFTER) {
72            return;
73        }
74        *last = Some(Instant::now());
75    }
76    let url = format!("{}/agentmux/agents/self/keys", local_url.trim_end_matches('/'));
77    let Ok(resp) = client.get(&url).header(AUTH_KEY_HEADER, auth_key).timeout(FETCH_TIMEOUT).send().await else { return };
78    if !resp.status().is_success() {
79        return;
80    }
81    if let Ok(keys) = resp.json::<Wire>().await {
82        *CACHE.lock().unwrap_or_else(|e| e.into_inner()) = Some(Cached { keys, fetched_at: Instant::now() });
83    }
84}
85
86/// Whether a cached copy needs no re-fetch at `now` (unix seconds).
87fn still_good(c: &Cached, now: i64) -> bool {
88    c.fetched_at.elapsed() < REFRESH_AFTER && c.keys.jekt_key_expires_at.is_none_or(|expires| now < expires)
89}
90
91/// A fetched name-keyed key, decoded, when srv served it for exactly the name
92/// this process signs as.
93fn name_key(pick: impl Fn(&Wire) -> Option<&String>) -> Option<Vec<u8>> {
94    let own = std::env::var("AGENTMUX_AGENT_ID").ok()?;
95    let cache = CACHE.lock().unwrap_or_else(|e| e.into_inner());
96    let cached = cache.as_ref()?;
97    if cached.keys.slug != own {
98        return None;
99    }
100    agentmux_common::jekt_sign::decode_key(pick(&cached.keys)?)
101}
102
103/// The host jekt HMAC key: the fetched one when it is ours, else the env's.
104pub(crate) fn jekt_key() -> Option<Vec<u8>> {
105    name_key(|k| k.jekt_key.as_ref()).or_else(|| env_key("AGENTMUX_JEKT_KEY"))
106}
107
108/// The LAN (and cross-channel) private key: fetched when ours, else the env's.
109pub(crate) fn lan_key() -> Option<Vec<u8>> {
110    name_key(|k| k.lan_key.as_ref()).or_else(|| env_key("AGENTMUX_LAN_KEY"))
111}
112
113/// The WAN private key: fetched when ours, else the env's.
114pub(crate) fn wan_key() -> Option<Vec<u8>> {
115    name_key(|k| k.wan_key.as_ref()).or_else(|| env_key("AGENTMUX_WAN_KEY"))
116}
117
118/// This agent's UID and its UID-keyed LAN private key, for the v2 signature
119/// (identity M4d-6). Only ever the fetched key, never an env one: the v2
120/// signature claims the UID, so it must come from srv's own record of it.
121pub(crate) fn uid_signer() -> Option<(String, Vec<u8>)> {
122    let cache = CACHE.lock().unwrap_or_else(|e| e.into_inner());
123    let keys = &cache.as_ref()?.keys;
124    let uid = keys.uid.trim();
125    if uid.is_empty() {
126        return None;
127    }
128    Some((uid.to_string(), agentmux_common::jekt_sign::decode_key(keys.uid_lan_key.as_ref()?)?))
129}
130
131fn env_key(var: &str) -> Option<Vec<u8>> {
132    std::env::var(var).ok().filter(|s| !s.is_empty()).and_then(|b64| agentmux_common::jekt_sign::decode_key(&b64))
133}
134
135#[cfg(test)]
136pub(crate) fn set_for_test(slug: &str, jekt_key: Option<&str>) {
137    *CACHE.lock().unwrap_or_else(|e| e.into_inner()) = Some(Cached {
138        keys: Wire { slug: slug.to_string(), jekt_key: jekt_key.map(str::to_string), ..Default::default() },
139        fetched_at: Instant::now(),
140    });
141}
142
143#[cfg(test)]
144pub(crate) fn clear_for_test() {
145    *CACHE.lock().unwrap_or_else(|e| e.into_inner()) = None;
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    /// [7; 32] and [1; 32], base64.
153    const FETCHED_B64: &str = "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=";
154    const ENV_B64: &str = "AQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQE=";
155
156    /// The cache and `AGENTMUX_AGENT_ID` / `AGENTMUX_JEKT_KEY` are process-wide.
157    static LOCK: Mutex<()> = Mutex::new(());
158
159    fn with_env<T>(agent_id: &str, env_key: Option<&str>, f: impl FnOnce() -> T) -> T {
160        let _guard = LOCK.lock().unwrap_or_else(|e| e.into_inner());
161        std::env::set_var("AGENTMUX_AGENT_ID", agent_id);
162        match env_key {
163            Some(k) => std::env::set_var("AGENTMUX_JEKT_KEY", k),
164            None => std::env::remove_var("AGENTMUX_JEKT_KEY"),
165        }
166        let out = f();
167        clear_for_test();
168        std::env::remove_var("AGENTMUX_JEKT_KEY");
169        out
170    }
171
172    #[test]
173    fn a_fetched_key_for_exactly_this_name_wins_over_the_env_copy() {
174        let fetched = [7u8; 32];
175        let key = with_env("aria", Some(ENV_B64), || {
176            set_for_test("aria", Some(FETCHED_B64));
177            jekt_key()
178        });
179        assert_eq!(key.as_deref(), Some(&fetched[..]));
180    }
181
182    #[test]
183    fn a_key_served_for_another_name_is_not_used() {
184        // This agent signs as "aria" but its row's slug is "aria-2": the
185        // served key is aria-2's, and srv checks "aria" against aria's.
186        let env = [1u8; 32];
187        let key = with_env("aria", Some(ENV_B64), || {
188            set_for_test("aria-2", Some(FETCHED_B64));
189            jekt_key()
190        });
191        assert_eq!(key.as_deref(), Some(&env[..]));
192    }
193
194    #[test]
195    fn nothing_fetched_means_the_env_key_as_before() {
196        let env = [1u8; 32];
197        assert_eq!(with_env("aria", Some(ENV_B64), jekt_key).as_deref(), Some(&env[..]));
198        assert_eq!(with_env("aria", None, jekt_key), None);
199    }
200
201    #[test]
202    fn a_copy_is_good_until_its_key_expires_then_refetched() {
203        let cached = |expires: Option<i64>| Cached {
204            keys: Wire { slug: "aria".into(), jekt_key_expires_at: expires, ..Default::default() },
205            fetched_at: Instant::now(),
206        };
207        assert!(still_good(&cached(Some(1_000)), 999));
208        assert!(!still_good(&cached(Some(1_000)), 1_000), "expired: srv may rotate it now");
209        assert!(still_good(&cached(None), 1_000_000), "no expiry served: the 20 h clock");
210    }
211
212    #[test]
213    fn the_v2_signer_is_only_ever_the_fetched_uid_key() {
214        let _guard = LOCK.lock().unwrap_or_else(|e| e.into_inner());
215        clear_for_test();
216        assert!(
217            uid_signer().is_none(),
218            "no fetch: no v2, whatever the env holds"
219        );
220        let set = |uid: &str, key: Option<&str>| {
221            *CACHE.lock().unwrap_or_else(|e| e.into_inner()) = Some(Cached {
222                keys: Wire {
223                    slug: "aria".into(),
224                    uid: uid.into(),
225                    uid_lan_key: key.map(str::to_string),
226                    ..Default::default()
227                },
228                fetched_at: Instant::now(),
229            });
230        };
231        set("uid-1", Some(FETCHED_B64));
232        assert_eq!(uid_signer(), Some(("uid-1".to_string(), vec![7u8; 32])));
233        set("", Some(FETCHED_B64));
234        assert!(uid_signer().is_none(), "a key with no UID to claim");
235        set("uid-1", None);
236        assert!(uid_signer().is_none());
237        clear_for_test();
238    }
239}