Skip to main content

agentmux_srv/server/
caller.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Who is calling, derived from the request's credential — identity M4a
5//! (`docs/specs/SPEC_AGENT_IDENTITY_CARRIED_NOT_DERIVED_2026_09_23.md`
6//! §6.5.3).
7//!
8//! A request that carries `X-Agent-Token` for a token this srv minted is
9//! [`Caller::Agent`] with that token's UID. Anything else is
10//! [`Caller::Unattributed`] — the UI, `bashwrap`, `muxsh`, a tokenless agent,
11//! a peer srv. **This is attribution, never authorization:** at the same-user
12//! boundary a request that omits its token cannot be told apart from the UI
13//! (§6.5.1), so nothing here refuses a request. An unknown or foreign-channel
14//! token is counted and treated as absent — an agent deleted while running,
15//! or a token that reached another channel's srv through a forwarded
16//! environment, keeps working exactly as before.
17//!
18//! Derived only for requests that authenticated with the full instance key:
19//! a token on a LAN-key request must not lift it to host trust (§6.5.3).
20
21use axum::body::Body;
22use axum::extract::State;
23use axum::http::Request;
24use axum::middleware::Next;
25use axum::response::Response;
26
27use super::{AppState, ReactiveAuthVia};
28
29/// The request header an agent's MCP sets from its inherited
30/// `AGENTMUX_AGENT_TOKEN`.
31pub(crate) const AGENT_TOKEN_HEADER: &str = "X-Agent-Token";
32
33/// Who a request is attributed to. Inserted into every request's extensions
34/// on the authed and LAN-forward route groups.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub(crate) enum Caller {
37    /// Carried a token this srv minted for `uid`.
38    Agent { uid: String },
39    /// Anything else. Not an error, and not a lesser principal: most callers
40    /// (the UI among them) are this.
41    Unattributed,
42}
43
44impl Caller {
45    /// The caller's UID, if attributed.
46    pub(crate) fn uid(&self) -> Option<&str> {
47        match self {
48            Caller::Agent { uid } => Some(uid),
49            Caller::Unattributed => None,
50        }
51    }
52}
53
54/// The UID to persist beside an actor name (identity M4c-1, spec §6.5.9):
55/// the request's attributed UID, or `""` — unknown, never guessed — when the
56/// request is Unattributed or carries no `Caller` at all.
57pub(crate) fn attributed_uid(caller: Option<&Caller>) -> String {
58    caller
59        .and_then(Caller::uid)
60        .map(str::to_string)
61        .unwrap_or_default()
62}
63
64/// Resolve the request's [`Caller`]. Pure, so the rules are testable without
65/// a router: `token` is the header value, `full_key` whether the request
66/// authenticated with the full instance key, `lookup` the token index.
67/// What the token index said about a token.
68pub(crate) enum Lookup {
69    Index(Option<String>),
70    /// No index attached (it failed to build at boot) — counted apart from an
71    /// unknown token, so the two are not confused.
72    NoIndex,
73}
74
75pub(crate) fn derive_caller(
76    token: Option<&str>,
77    full_key: bool,
78    lookup: impl FnOnce(&str) -> Lookup,
79) -> Caller {
80    let Some(token) = token.map(str::trim).filter(|t| !t.is_empty()) else {
81        return Caller::Unattributed;
82    };
83    if !full_key {
84        crate::backend::agent_resolve::record_uid_fallback("m4.token_on_lan_key");
85        return Caller::Unattributed;
86    }
87    match lookup(token) {
88        Lookup::Index(Some(uid)) => Caller::Agent { uid },
89        Lookup::Index(None) => {
90            crate::backend::agent_resolve::record_uid_fallback("m4.token_unknown");
91            Caller::Unattributed
92        }
93        Lookup::NoIndex => {
94            crate::backend::agent_resolve::record_uid_fallback("m4.token_index_unavailable");
95            Caller::Unattributed
96        }
97    }
98}
99
100/// Inserts [`Caller`]. Layered *inside* `auth_middleware` /
101/// `lan_or_full_auth_middleware`; both mark which key authenticated the
102/// request (`ReactiveAuthVia`), and a `Caller` is derived only when that mark
103/// says the full key — a request that reached here unmarked is Unattributed,
104/// never assumed full-key (review on #3571). The WebSocket upgrade is always
105/// Unattributed (§6.5.3): WebSocket RPC is the UI's channel, and a token
106/// must never ride on it.
107pub(crate) async fn caller_middleware(
108    State(state): State<AppState>,
109    mut req: Request<Body>,
110    next: Next,
111) -> Response {
112    let full_key = req.extensions().get::<ReactiveAuthVia>() == Some(&ReactiveAuthVia::FullAuthKey);
113    // A CORS preflight skips auth, so it is never marked; it is not a
114    // token on a LAN key and must not be counted as one.
115    // A container token names its agent itself; a header can't re-point it.
116    let pinned = req
117        .extensions()
118        .get::<crate::backend::container_credential::ContainerGrant>()
119        .map(|g| g.agent_token.clone());
120    let header = if req.uri().path() == "/ws" || req.method() == axum::http::Method::OPTIONS {
121        None
122    } else {
123        req.headers()
124            .get(AGENT_TOKEN_HEADER)
125            .and_then(|v| v.to_str().ok())
126    };
127    let token = match &pinned {
128        Some(agent_token) => agent_token.as_deref(),
129        None => header,
130    };
131    let caller = derive_caller(token, full_key, |t| match state.mstore.token_index() {
132        Some(index) => Lookup::Index(index.uid_for(t)),
133        None => Lookup::NoIndex,
134    });
135    req.extensions_mut().insert(caller);
136    next.run(req).await
137}
138
139/// `GET /agentmux/identity/fallbacks` — every §9.2 counter and the two
140/// live gauges (§6.5.6). The phase gates read this: M4d and M5 wait for
141/// `live.tokenless_or_unknown` to reach zero and for the fallback counters
142/// to stop moving. Read-only; auth-gated like every loopback route.
143///
144/// Also reports who the request itself was attributed to (`caller`), so an
145/// agent can check that its token is being recognised.
146pub(crate) async fn handle_identity_fallbacks(
147    State(state): State<AppState>,
148    caller: Option<axum::Extension<Caller>>,
149) -> axum::Json<serde_json::Value> {
150    let counters: serde_json::Map<String, serde_json::Value> =
151        crate::backend::agent_resolve::uid_fallback_counts()
152            .into_iter()
153            .map(|(site, n)| (site.to_string(), serde_json::json!(n)))
154            .collect();
155    // `has_row` reads the store per UID-less registration, so off the async
156    // worker (incident #1782).
157    let registrations = state.reactive_handler.list_agents();
158    let mstore = state.mstore.clone();
159    // A failed read reports the gauges as null, never as a drained zero.
160    let gauges = tokio::task::spawn_blocking(move || {
161        crate::backend::identity_spawn::live_gauges(
162            &registrations,
163            |block| crate::backend::blockcontroller::get_controller(block).is_some(),
164            |block| crate::backend::identity_spawn::store_block_has_row(&mstore, block),
165        )
166    })
167    .await
168    .ok();
169    let (tokenless_or_unknown, unidentified) = match gauges {
170        Some((t, u)) => (serde_json::json!(t), serde_json::json!(u)),
171        None => (serde_json::Value::Null, serde_json::Value::Null),
172    };
173    let caller_uid = caller.as_ref().and_then(|c| c.0.uid().map(str::to_string));
174    axum::Json(serde_json::json!({
175        // Counters and gauges are in memory, per srv, since this boot: a
176        // counter that has not fired is absent (read it as zero), and a gate
177        // reader must sample every channel's srv (spec §9.2).
178        "boot_id": state.boot_id.to_string(),
179        "caller": caller_uid,
180        "counters": counters,
181        "gauges": {
182            "live.tokenless_or_unknown": tokenless_or_unknown,
183            "live.unidentified": unidentified,
184        },
185    }))
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    fn index(token: &str) -> Lookup {
193        Lookup::Index((token == "tok-y").then(|| "uid-y".to_string()))
194    }
195
196    #[test]
197    fn a_known_token_on_the_full_key_is_the_agent() {
198        assert_eq!(
199            derive_caller(Some("tok-y"), true, index),
200            Caller::Agent {
201                uid: "uid-y".into()
202            }
203        );
204    }
205
206    #[test]
207    fn no_token_is_unattributed() {
208        assert_eq!(derive_caller(None, true, index), Caller::Unattributed);
209        assert_eq!(derive_caller(Some("  "), true, index), Caller::Unattributed);
210    }
211
212    /// Never a 401: an unknown token (deleted agent, another channel's srv)
213    /// is treated as absent.
214    #[test]
215    fn an_unknown_token_is_unattributed_not_refused() {
216        assert_eq!(
217            derive_caller(Some("tok-gone"), true, index),
218            Caller::Unattributed
219        );
220    }
221
222    #[test]
223    fn a_missing_index_is_unattributed() {
224        assert_eq!(
225            derive_caller(Some("tok-y"), true, |_| Lookup::NoIndex),
226            Caller::Unattributed
227        );
228    }
229
230    /// A token must not lift a LAN-key request to host trust.
231    #[test]
232    fn a_token_on_a_lan_key_request_is_ignored() {
233        assert_eq!(
234            derive_caller(Some("tok-y"), false, index),
235            Caller::Unattributed
236        );
237    }
238}