agentmux_srv/server/app_api/
blockstats.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! `blockstats.detail`: a terminal's stats badge and its tooltip
5//! (`backend::blockstats_detail`). The window's, not an agent API: a
6//! connection registered as an agent is refused.
7//!
8//! - A local pane: its processes here, from one snapshot.
9//! - A remote pane (SSH, WSL): its processes on that machine, asked of
10//!   AgentMux's helper there, and the machine's name and address.
11
12use super::*;
13use super::tower_pane::not_an_agent_api;
14use crate::backend::blockcontroller::durable_ssh::META_KEY_SESSION_ID;
15use crate::backend::blockcontroller::process_tree::MAX_PIDS_PER_BLOCK;
16use crate::backend::blockstats_detail;
17use crate::backend::remote::address as remote_address;
18use crate::backend::remote::ConnTarget;
19
20/// How long asking a remote pane's helper may take before the badge reads
21/// "not measured" for this round.
22const REMOTE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
23
24pub fn register(engine: &Arc<WshRpcEngine>, state: &AppState) {
25    let mstore = state.mstore.clone();
26    engine.register_typed(COMMAND_BLOCKSTATS_DETAIL, move |req: BlockStatsDetailReq, ctx| {
27        let mstore = mstore.clone();
28        async move {
29            not_an_agent_api(&ctx)?;
30            let meta = mstore.get::<Block>(&req.block_id).ok().flatten().map(|b| b.meta);
31            let meta_str = |key: &str| meta.as_ref().map(|m| obj::meta_get_string(m, key, "")).unwrap_or_default();
32            let conn = meta_str("connection");
33            // Only a terminal has a remote shell to measure; any other pane on
34            // a connection (Files, preview) is not asked of the machine.
35            let remote_terminal = meta_str("view") == "term";
36            if remote_terminal && matches!(ConnTarget::parse(&conn), Ok(ConnTarget::Ssh(_) | ConnTarget::Wsl(_))) {
37                // The machine's name and address, and the pane's processes
38                // there, side by side.
39                let session = meta_str(META_KEY_SESSION_ID);
40                let (connection, remote) = tokio::join!(connection(&conn), remote(&conn, &req.block_id, &session));
41                return Ok(BlockStatsDetail { local: None, connection, remote: Some(remote) });
42            }
43            let root = blockcontroller::pidregistry::get_all()
44                .into_iter()
45                .find(|(id, _)| *id == req.block_id)
46                .map(|(_, pid)| pid)
47                .ok_or("no process for this pane")?;
48            // The process table: off the async workers.
49            let local = tokio::task::spawn_blocking(move || {
50                let snap = agentmux_procstats::snapshot().map_err(|e| format!("blockstats.detail: {e}"))?;
51                blockstats_detail::detail(&snap, root, MAX_PIDS_PER_BLOCK)
52                    .ok_or_else(|| "the pane's shell has exited".to_string())
53            })
54            .await
55            .map_err(|e| format!("blockstats.detail: {e}"))??;
56            Ok(BlockStatsDetail { local: Some(local), connection: None, remote: None })
57        }
58    });
59}
60
61/// A remote pane's machine; `None` for a local pane.
62async fn connection(name: &str) -> Option<BlockStatsConnection> {
63    let kind = match ConnTarget::parse(name).ok()? {
64        ConnTarget::Local => return None,
65        ConnTarget::Ssh(_) => "ssh",
66        ConnTarget::Wsl(_) => "wsl",
67    };
68    let found = remote_address::of(name).await.unwrap_or_default();
69    Some(BlockStatsConnection {
70        name: name.trim().to_string(),
71        kind: kind.to_string(),
72        host: found.host,
73        address: found.address,
74    })
75}
76
77/// A remote pane's processes on its machine. WSL: the helper Tower runs in
78/// the distribution. SSH: the host's helper, only over a connection that
79/// needs no prompt and no install (`files::connect_quiet`): this is asked
80/// every few seconds, never on the user's word.
81async fn remote(conn: &str, block_id: &str, session: &str) -> BlockStatsRemote {
82    let ask = async {
83        let helper = match ConnTarget::parse(conn)? {
84            ConnTarget::Wsl(distro) => crate::backend::tower_remote::wsl_helper(&distro).await?,
85            ConnTarget::Ssh(_) => crate::backend::remote::files::connect_quiet(conn).await.map_err(|e| e.message)?,
86            ConnTarget::Local => return Err("not a remote pane".to_string()),
87        };
88        helper.pane_procs(block_id, session).await.map_err(|e| e.message)
89    };
90    match tokio::time::timeout(REMOTE_TIMEOUT, ask).await {
91        Ok(Ok(frame)) => blockstats_detail::remote(block_id, frame),
92        Ok(Err(why)) => blockstats_detail::remote_unavailable(why),
93        Err(_) => blockstats_detail::remote_unavailable("the machine didn't answer in time".to_string()),
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    async fn call(data: serde_json::Value, agent_id: &str) -> (serde_json::Value, String) {
102        let state = crate::server::tests::test_state();
103        let (engine, mut rx) = WshRpcEngine::new();
104        register(&engine, &state);
105        if !agent_id.is_empty() {
106            engine.set_rpc_context(RpcContext { agent_id: agent_id.to_string(), ..Default::default() });
107        }
108        engine.handle_message(RpcMessage {
109            command: COMMAND_BLOCKSTATS_DETAIL.to_string(),
110            reqid: "req-1".to_string(),
111            data: Some(data),
112            ..Default::default()
113        });
114        let resp = tokio::time::timeout(std::time::Duration::from_secs(20), rx.recv()).await.unwrap().unwrap();
115        (resp.data.unwrap_or_default(), resp.error)
116    }
117
118    #[tokio::test(flavor = "multi_thread")]
119    async fn a_registered_pane_gets_its_process_tree() {
120        // This test process stands in for the pane's shell.
121        let block_id = "blockstats-detail-test-block";
122        blockcontroller::pidregistry::register(block_id, std::process::id());
123        let (data, err) = call(json!({ "block_id": block_id }), "").await;
124        blockcontroller::pidregistry::unregister(block_id);
125        assert_eq!(err, "");
126        let d: BlockStatsDetail = serde_json::from_value(data).unwrap();
127        assert!(d.connection.is_none() && d.remote.is_none(), "{d:?}");
128        let d = d.local.expect("a local pane's figures");
129        assert!(d.processes >= 1);
130        assert!(d.started_at_ms.is_some());
131        assert!(d.mem_private.unwrap_or(0) > 0, "{d:?}");
132        assert_eq!(d.mem_commit.is_some(), cfg!(windows), "{d:?}");
133    }
134
135    #[tokio::test(flavor = "multi_thread")]
136    async fn an_unknown_pane_is_an_error() {
137        let (_, err) = call(json!({ "block_id": "no-such-block" }), "").await;
138        assert!(err.contains("no process"), "{err}");
139    }
140
141    #[tokio::test(flavor = "multi_thread")]
142    async fn a_local_pane_has_no_connection_and_a_wsl_pane_names_its_distro() {
143        assert_eq!(connection("").await, None);
144        assert_eq!(connection("local").await, None);
145        let wsl = connection("wsl://NoSuchDistro").await.unwrap();
146        assert_eq!((wsl.kind.as_str(), wsl.host.as_deref()), ("wsl", Some("NoSuchDistro")));
147        assert_eq!(wsl.address, None, "a distro that doesn't exist has no address");
148    }
149
150    /// An SSH pane's machine answers through its helper's connection (here
151    /// the real request handling, in process): this test process carries no
152    /// pane's marker, so the host's figures come back without the pane's.
153    #[tokio::test(flavor = "multi_thread")]
154    async fn a_remote_pane_is_measured_by_its_machines_helper() {
155        let home = tempfile::tempdir().unwrap();
156        crate::backend::remote::files::connect_in_process("me@stats-test-host", home.path().to_path_buf()).await;
157        let r = remote("me@stats-test-host", "no-pane-has-this-id", "").await;
158        assert_eq!(r.found, "none", "{r:?}");
159        let host = r.host.expect("the host's figures");
160        assert!(host.processes > 1 && host.cpu_count >= 1, "{host:?}");
161        assert!(r.figures.is_none());
162    }
163
164    #[tokio::test(flavor = "multi_thread")]
165    async fn a_remote_pane_whose_helper_isnt_there_is_unavailable_not_an_error() {
166        let r = remote("me@no-helper-here.invalid", "b", "").await;
167        assert_eq!(r.found, "unavailable", "{r:?}");
168        assert!(r.reason.is_some());
169    }
170
171    #[tokio::test(flavor = "multi_thread")]
172    async fn a_registered_agent_is_refused() {
173        let (_, err) = call(json!({ "block_id": "x" }), "AgentX").await;
174        assert!(err.contains("FORBIDDEN"), "{err}");
175    }
176}
agentmux_srv/server/app_api/blockstats.rs

agentmux_srv/server/app_api/
blockstats.rs
