objectiveai_sdk/laboratories/filetree/client.rs
1//! Materialized consumer of a live file-tree SSE stream — a
2//! laboratory MCP's own `/filetree` endpoint
3//! ([`FileTree::laboratory`]) or a CLI daemon's
4//! `/laboratories/{id}/filetree` proxy ([`FileTree::daemon`]); both
5//! speak the identical wire contract.
6//!
7//! [`FileTree`] is NOT a raw event stream — it connects once, then
8//! folds every incoming [`FileTreeEvent`] into an in-memory,
9//! self-updating recursive tree (the watched root's
10//! [`FileTreeNode`](super::FileTreeNode) children): a
11//! [`Snapshot`](FileTreeEvent::Snapshot) replaces the whole child set,
12//! [`Upserted`](FileTreeEvent::Upserted) inserts/replaces one node at
13//! its path (a directory node carries its whole subtree), and
14//! [`Removed`](FileTreeEvent::Removed) drops the node at its path.
15//!
16//! Two ways to observe it:
17//! - [`tree`](FileTree::tree) — async snapshot of the current tree.
18//! - an on-change **callback**
19//! ([`on_change`](FileTreeBuilder::on_change)), invoked with the full
20//! refreshed tree on every applied change.
21//! - [`subscribe`](FileTree::subscribe) — async, blocks until the next
22//! change.
23//!
24//! One [`FileTree`] = one connection: the internal pump runs until the
25//! laboratory closes the stream; after that the view is frozen at its
26//! last state. Dropping it aborts the pump. Reconnection is the
27//! caller's loop — build a new [`FileTree`].
28
29use std::sync::Arc;
30
31use futures::StreamExt;
32use reqwest_eventsource::{Event, RequestBuilderExt};
33use tokio::sync::{Mutex, watch};
34
35use super::{FileTreeEvent, FileTreeNode};
36
37/// The on-change callback: invoked with the full current tree (the
38/// watched root's child nodes) after each applied [`FileTreeEvent`].
39pub type OnChange = Box<dyn Fn(&[FileTreeNode]) + Send + Sync>;
40
41#[derive(Debug, thiserror::Error)]
42pub enum Error {
43 /// The request builder rejected the URL, or opening the SSE
44 /// stream failed.
45 #[error("connect laboratory filetree: {0}")]
46 Connect(#[from] reqwest_eventsource::CannotCloneRequestError),
47 /// The underlying reqwest client failed to build.
48 #[error("laboratory filetree http client: {0}")]
49 Client(#[from] reqwest::Error),
50}
51
52/// The shared inner state, held by both the [`FileTree`] handle and
53/// its pump task.
54struct Shared {
55 /// The tree itself: the watched root's child nodes, kept live by
56 /// folding each event.
57 state: Mutex<Vec<FileTreeNode>>,
58 /// A monotonically-bumped change counter. Each applied event bumps
59 /// it, waking every [`subscribe`](FileTree::subscribe) waiter.
60 changes: watch::Sender<u64>,
61 /// Optional push callback, invoked with the full tree after each
62 /// change.
63 on_change: Option<OnChange>,
64}
65
66/// Unconnected configuration — [`FileTree::laboratory`] /
67/// [`FileTree::daemon`] + builder methods + [`FileTreeBuilder::connect`].
68pub struct FileTreeBuilder {
69 /// The full endpoint URL — `{lab}/filetree` (laboratory variant) or
70 /// `{daemon}/laboratories/{id}/filetree` (daemon variant).
71 url: String,
72 /// Query params to send (the laboratory variant's `path` scope, the
73 /// daemon variant's host pin).
74 query: Vec<(String, String)>,
75 /// Optional daemon auth signature, sent as the
76 /// `X-OBJECTIVEAI-SIGNATURE` header (daemon variant only —
77 /// laboratories have no auth).
78 signature: Option<String>,
79 /// Optional on-change callback.
80 on_change: Option<OnChange>,
81}
82
83impl FileTreeBuilder {
84 /// Scope the watch to an absolute path inside the container.
85 /// Laboratory variant only: the daemon endpoint always serves the
86 /// laboratory's workspace (its cwd) and ignores this.
87 pub fn path(mut self, path: impl Into<String>) -> Self {
88 self.query.push(("path".to_string(), path.into()));
89 self
90 }
91
92 /// Attach the daemon auth signature (the pre-derived
93 /// `sha256=<hex(SHA256(DAEMON_SECRET))>`), sent as the
94 /// `X-OBJECTIVEAI-SIGNATURE` header. Daemon variant only; without
95 /// it the daemon must be running without a secret.
96 pub fn signature(mut self, signature: impl Into<String>) -> Self {
97 self.signature = Some(signature.into());
98 self
99 }
100
101 /// Pin the exact serving host (`?machine=…&machine_state=…`) —
102 /// laboratory ids are only unique per (machine, state). Daemon
103 /// variant only; omit for the legacy first-match-by-id scan.
104 pub fn host(
105 mut self,
106 machine: impl Into<String>,
107 machine_state: impl Into<String>,
108 ) -> Self {
109 self.query.push(("machine".to_string(), machine.into()));
110 self
111 .query
112 .push(("machine_state".to_string(), machine_state.into()));
113 self
114 }
115
116 /// Register a callback invoked with the full current tree after
117 /// every applied change. Runs on the pump task, so keep it cheap
118 /// and non-blocking; for the tree on demand use
119 /// [`tree`](FileTree::tree).
120 pub fn on_change(
121 mut self,
122 callback: impl Fn(&[FileTreeNode]) + Send + Sync + 'static,
123 ) -> Self {
124 self.on_change = Some(Box::new(callback));
125 self
126 }
127
128 /// Open the SSE stream and start the pump. The returned
129 /// [`FileTree`] immediately begins folding events (the first is
130 /// the endpoint's connect-time snapshot).
131 pub async fn connect(self) -> Result<FileTree, Error> {
132 let client = reqwest::Client::builder().build()?;
133 let mut request = client
134 .get(&self.url)
135 .header("Accept", "text/event-stream");
136 if !self.query.is_empty() {
137 request = request.query(&self.query);
138 }
139 if let Some(signature) = &self.signature {
140 request = request.header("X-OBJECTIVEAI-SIGNATURE", signature);
141 }
142 let source = request.eventsource()?;
143
144 let shared = Arc::new(Shared {
145 state: Mutex::new(Vec::new()),
146 changes: watch::channel(0u64).0,
147 on_change: self.on_change,
148 });
149 let pump = tokio::spawn(pump(source, shared.clone()));
150 Ok(FileTree { shared, pump })
151 }
152}
153
154/// The materialized `/filetree` view — see the module docs. Construct
155/// via [`FileTree::laboratory`] or [`FileTree::daemon`]. Dropping it
156/// aborts the background pump.
157pub struct FileTree {
158 shared: Arc<Shared>,
159 pump: tokio::task::JoinHandle<()>,
160}
161
162impl FileTree {
163 /// Build a [`FileTree`] against a laboratory MCP's own `/filetree`
164 /// endpoint, from its base HTTP address (e.g.
165 /// `http://127.0.0.1:<port>` — the loopback-published container
166 /// port). No auth — the port itself is the trust boundary.
167 pub fn laboratory(base_url: impl Into<String>) -> FileTreeBuilder {
168 FileTreeBuilder {
169 url: format!("{}/filetree", base_url.into()),
170 query: Vec::new(),
171 signature: None,
172 on_change: None,
173 }
174 }
175
176 /// Build a [`FileTree`] against a CLI daemon's
177 /// `/laboratories/{id}/filetree` endpoint, from the daemon's
178 /// published `http://` address and the RAW laboratory id (ids
179 /// forbid `/`, so the URL is safe by construction). Same wire
180 /// contract as the laboratory variant — the daemon re-emits the
181 /// identical events from its own materialized state. Attach
182 /// [`signature`](FileTreeBuilder::signature) against a secretful
183 /// daemon and [`host`](FileTreeBuilder::host) to pin the exact
184 /// serving host.
185 pub fn daemon(
186 base_url: impl Into<String>,
187 laboratory_id: impl AsRef<str>,
188 ) -> FileTreeBuilder {
189 FileTreeBuilder {
190 url: format!(
191 "{}/laboratories/{}/filetree",
192 base_url.into(),
193 laboratory_id.as_ref(),
194 ),
195 query: Vec::new(),
196 signature: None,
197 on_change: None,
198 }
199 }
200
201 /// Snapshot the current tree — the watched root's child nodes.
202 /// Empty before the first snapshot has been folded.
203 pub async fn tree(&self) -> Vec<FileTreeNode> {
204 self.shared.state.lock().await.clone()
205 }
206
207 /// Block until the next change is applied. A fresh call waits for
208 /// the FIRST change after it is made, so a change that lands
209 /// between a preceding [`tree`](Self::tree) read and this
210 /// call is not observed by it — pair with the read in a loop, or
211 /// use the [`on_change`](FileTreeBuilder::on_change) callback for
212 /// guaranteed push.
213 pub async fn subscribe(&self) {
214 let mut rx = self.shared.changes.subscribe();
215 let _ = rx.changed().await;
216 }
217}
218
219impl Drop for FileTree {
220 fn drop(&mut self) {
221 self.pump.abort();
222 }
223}
224
225/// Read SSE messages, fold each [`FileTreeEvent`] into `shared.state`,
226/// fire the callback with the refreshed set, and bump the change
227/// counter. Runs until the stream closes. Parse errors and non-message
228/// events are skipped; a transport error ends the pump (the caller
229/// rebuilds — reconnection is not automatic).
230async fn pump(mut source: reqwest_eventsource::EventSource, shared: Arc<Shared>) {
231 while let Some(event) = source.next().await {
232 match event {
233 Ok(Event::Open) => continue,
234 Ok(Event::Message(message)) => {
235 let Ok(event) = serde_json::from_str::<FileTreeEvent>(&message.data)
236 else {
237 // Skip a frame we can't parse rather than tearing down.
238 continue;
239 };
240 let snapshot = {
241 let mut state = shared.state.lock().await;
242 event.apply(&mut state);
243 shared.on_change.as_ref().map(|_| state.clone())
244 };
245 if let (Some(callback), Some(snapshot)) =
246 (&shared.on_change, &snapshot)
247 {
248 callback(snapshot);
249 }
250 shared.changes.send_modify(|version| {
251 *version = version.wrapping_add(1);
252 });
253 }
254 // Transport error (including the stream closing): stop.
255 Err(_) => break,
256 }
257 }
258}