okbot_peer/
node.rs

1//! A bot's place in its teams: manifests, keys, and the transport to other members.
2
3use std::collections::{BTreeMap, HashMap, VecDeque};
4use std::path::{Path, PathBuf};
5use std::sync::{Arc, Mutex, OnceLock, RwLock};
6use std::time::{Duration, Instant};
7
8use jiff::Timestamp;
9use okbot_core::{
10    Clock, ConversationKey, ConversationStore, InboundSink, PeerCaller, PeerError, Prompt, SentRef,
11};
12use serde::{Deserialize, Serialize};
13use serde_json::Value;
14use sha2::Digest as _;
15use tokio::io::{AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt as _, BufReader};
16
17use crate::key::{PublicKey, SigningKey, bot_public_key};
18use crate::manifest::{self, Manifest, Member};
19use crate::token::{self, Claims, Replays};
20use crate::{Error, a2a};
21
22/// Where the sockets of the bots of this machine live: `$OKBOT_PEER_DIR`, else
23/// `$XDG_RUNTIME_DIR/okbot`, else `<temp>/okbot-<user>`. Bots under different OS users share one
24/// directory by setting `OKBOT_PEER_DIR` alike.
25pub fn socket_dir() -> PathBuf {
26    if let Some(dir) = std::env::var_os("OKBOT_PEER_DIR") {
27        return PathBuf::from(dir);
28    }
29    if let Some(dir) = std::env::var_os("XDG_RUNTIME_DIR") {
30        return PathBuf::from(dir).join("okbot");
31    }
32    let user = std::env::var("USER").unwrap_or_else(|_| "user".into());
33    std::env::temp_dir().join(format!("okbot-{user}"))
34}
35
36/// The socket in `sockets` of the bot in `bot_dir`, named after its canonical path (short whatever
37/// the path's length, and the same for the bot and for the members that call it).
38pub fn socket_path(sockets: &Path, bot_dir: &Path) -> PathBuf {
39    let dir = std::fs::canonicalize(bot_dir).unwrap_or_else(|_| bot_dir.to_path_buf());
40    let digest = sha2::Sha256::digest(dir.as_os_str().as_encoded_bytes());
41    sockets.join(format!("{}.sock", &hex::encode(digest)[..20]))
42}
43/// Header carrying the response's signature.
44pub const SIGNATURE_HEADER: &str = "okbot-peer-signature";
45/// The authorization scheme.
46pub(crate) const SCHEME: &str = "OkbotPeer ";
47/// How often fetched manifests are refreshed.
48pub(crate) const REFRESH: Duration = Duration::from_secs(600);
49/// How long a member that did not answer is skipped.
50const DOWN_FOR: Duration = Duration::from_secs(60);
51/// Largest request or response body.
52pub(crate) const MAX_BODY: usize = 4 << 20;
53/// Time for one request to another member.
54const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
55
56/// Where a team's manifest comes from.
57#[derive(Debug, Clone, PartialEq, Eq)]
58pub enum ManifestSource {
59    /// A file on this machine (trusted as is).
60    File(PathBuf),
61    /// Built from `[a2a]` (the team `external`).
62    Builtin,
63    /// An `https://` URL, with `<url>.sig` signed by the team key.
64    Url {
65        /// The manifest URL.
66        url: String,
67        /// The pinned team key.
68        team_key: PublicKey,
69    },
70}
71
72/// This bot's settings for one team (`[teams.<id>]`).
73#[derive(Debug, Clone)]
74pub struct TeamConfig {
75    /// Team id.
76    pub id: String,
77    /// This bot's member name.
78    pub member: String,
79    /// Where the manifest is.
80    pub source: ManifestSource,
81    /// Other members may delegate to this bot.
82    pub accept: bool,
83    /// Longest chain accepted.
84    pub max_depth: u32,
85    /// Inbound tasks at once per caller.
86    pub max_concurrent: u32,
87    /// The longest a task delegated by this bot may take.
88    pub task_timeout: Duration,
89}
90
91/// Everything a [`Node`](crate::Node) needs besides its ports.
92#[derive(Debug, Clone)]
93pub struct NodeConfig {
94    /// This bot's key.
95    pub key: SigningKey,
96    /// This bot's X25519 key, for sealed messages through relays.
97    pub box_key: crate::seal::BoxKey,
98    /// This bot's `state/` directory.
99    pub state_dir: PathBuf,
100    /// Where the sockets of this machine's bots are ([`socket_dir`]).
101    pub sockets: PathBuf,
102    /// Teams.
103    pub teams: Vec<TeamConfig>,
104    /// A2A agents and clients that are not okbots, and this bot's name among them.
105    pub external: Option<(String, crate::external::ExternalConfig)>,
106}
107
108/// One team as this bot knows it.
109pub(crate) struct Team {
110    pub(crate) config: TeamConfig,
111    manifest: RwLock<Arc<Manifest>>,
112    /// When a request last made this bot read the manifest again.
113    forced: Mutex<Option<Instant>>,
114}
115
116impl Team {
117    pub(crate) fn manifest(&self) -> Arc<Manifest> {
118        self.manifest
119            .read()
120            .unwrap_or_else(std::sync::PoisonError::into_inner)
121            .clone()
122    }
123
124    pub(crate) fn signed(&self) -> bool {
125        matches!(self.config.source, ManifestSource::Url { .. })
126    }
127
128    /// Whether a request naming a newer manifest may make this bot read it now: at most every
129    /// [`FORCED_REFRESH`], since the claim is read before the signature can be checked.
130    pub(crate) fn may_force_refresh(&self) -> bool {
131        let mut last = lock(&self.forced);
132        if last.is_some_and(|at| at.elapsed() < FORCED_REFRESH) {
133            return false;
134        }
135        *last = Some(Instant::now());
136        true
137    }
138}
139
140/// The shortest time between two reads of a manifest asked for by requests.
141const FORCED_REFRESH: Duration = Duration::from_secs(30);
142
143/// An inbound task waiting for the runtime's answer.
144#[derive(Debug, Clone)]
145pub(crate) struct Running {
146    pub(crate) team: String,
147    pub(crate) caller: String,
148    pub(crate) task: String,
149    pub(crate) deadline: Timestamp,
150}
151
152/// A prompt the runtime showed in a peer conversation.
153#[derive(Debug, Clone)]
154pub(crate) struct OpenPrompt {
155    pub(crate) prompt: Prompt,
156    pub(crate) sent: SentRef,
157}
158
159/// Shared state of a node.
160pub(crate) struct Inner {
161    pub(crate) key: SigningKey,
162    pub(crate) box_key: crate::seal::BoxKey,
163    pub(crate) sockets: PathBuf,
164    pub(crate) state_dir: PathBuf,
165    pub(crate) teams: BTreeMap<String, Team>,
166    pub(crate) store: Arc<dyn ConversationStore>,
167    pub(crate) clock: Arc<dyn Clock>,
168    pub(crate) replays: Replays,
169    pub(crate) http: reqwest::Client,
170    pub(crate) sink: OnceLock<Arc<dyn InboundSink>>,
171    pub(crate) local_server: std::sync::atomic::AtomicBool,
172    /// Inbound tasks by conversation, oldest first.
173    pub(crate) inbound: Mutex<HashMap<ConversationKey, VecDeque<Running>>>,
174    pub(crate) callers: Mutex<HashMap<ConversationKey, PeerCaller>>,
175    /// Open prompts by inbound task id.
176    pub(crate) prompts: Mutex<HashMap<String, OpenPrompt>>,
177    /// Members that did not answer, until when.
178    down: Mutex<HashMap<(String, String), Instant>>,
179    /// Requests for members this bot relays.
180    pub(crate) relay: crate::relay::Relay,
181    /// Tokens reported for inbound tasks, until they end.
182    pub(crate) usage: Mutex<HashMap<String, u64>>,
183    /// `[a2a]` agents and clients.
184    pub(crate) external: Option<crate::external::ExternalConfig>,
185    /// What the turn of each working inbound task is doing.
186    pub(crate) progress: Mutex<HashMap<String, String>>,
187}
188
189pub(crate) fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
190    m.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
191}
192
193impl Inner {
194    pub(crate) fn new(
195        config: NodeConfig,
196        store: Arc<dyn ConversationStore>,
197        clock: Arc<dyn Clock>,
198    ) -> Result<Self, Error> {
199        let mut teams = BTreeMap::new();
200        for team in config.teams {
201            let manifest = match &team.source {
202                ManifestSource::File(path) => load_file(path, &team.id)?,
203                // Fetched before the node serves (see `Node::start`); empty until then.
204                ManifestSource::Url { .. } | ManifestSource::Builtin => Manifest::new(&team.id),
205            };
206            teams.insert(
207                team.id.clone(),
208                Team {
209                    config: team,
210                    manifest: RwLock::new(Arc::new(manifest)),
211                    forced: Mutex::new(None),
212                },
213            );
214        }
215        let external = config.external.map(|(me, external)| {
216            teams.insert(
217                crate::external::TEAM.to_string(),
218                Team {
219                    config: TeamConfig {
220                        id: crate::external::TEAM.into(),
221                        member: me.clone(),
222                        source: ManifestSource::Builtin,
223                        accept: !external.clients.is_empty(),
224                        max_depth: 1,
225                        max_concurrent: external.max_concurrent,
226                        task_timeout: external.task_timeout,
227                    },
228                    manifest: RwLock::new(Arc::new(external.manifest(&me))),
229                    forced: Mutex::new(None),
230                },
231            );
232            external
233        });
234        let _ = rustls::crypto::ring::default_provider().install_default();
235        Ok(Self {
236            key: config.key,
237            box_key: config.box_key,
238            sockets: config.sockets,
239            state_dir: config.state_dir,
240            teams,
241            store,
242            clock,
243            replays: Replays::default(),
244            http: reqwest::Client::builder()
245                .timeout(REQUEST_TIMEOUT)
246                .build()
247                .map_err(|e| Error::Key(format!("HTTP client: {e}")))?,
248            sink: OnceLock::new(),
249            local_server: std::sync::atomic::AtomicBool::new(false),
250            inbound: Mutex::default(),
251            callers: Mutex::default(),
252            prompts: Mutex::default(),
253            down: Mutex::default(),
254            relay: crate::relay::Relay::default(),
255            usage: Mutex::default(),
256            external,
257            progress: Mutex::default(),
258        })
259    }
260
261    pub(crate) fn now(&self) -> Timestamp {
262        self.clock.now()
263    }
264
265    /// A warning when a signed team's manifest expires within
266    /// [`manifest::EXPIRY_WARNING_DAYS`].
267    pub(crate) fn expiry_warning(&self, team: &Team) -> Option<String> {
268        let days = team
269            .manifest()
270            .days_left(self.now())
271            .filter(|d| *d < manifest::EXPIRY_WARNING_DAYS)?;
272        Some(format!(
273            "team `{}`: the manifest expires in {days} days; its owner signs it again \
274             (`okbot team renew {}`)",
275            team.config.id, team.config.id
276        ))
277    }
278
279    /// Checks that this bot is a member of each file-based team under its own key.
280    pub(crate) fn check_membership(&self) -> Vec<String> {
281        let mut problems = Vec::new();
282        for team in self.teams.values() {
283            let manifest = team.manifest();
284            if (team.signed() && manifest.team.version == 0)
285                || team.config.source == ManifestSource::Builtin
286            {
287                continue;
288            }
289            match manifest.members.get(&team.config.member) {
290                None => problems.push(format!(
291                    "team `{}`: this bot is not a member named `{}`",
292                    team.config.id, team.config.member
293                )),
294                Some(m) => {
295                    if member_key(m).ok().as_ref() != Some(&self.key.public()) {
296                        problems.push(format!(
297                            "team `{}`: member `{}` has another key than this bot",
298                            team.config.id, team.config.member
299                        ));
300                    }
301                }
302            }
303        }
304        problems
305    }
306
307    /// Replaces a team's manifest with a newer one.
308    pub(crate) fn install(&self, team: &Team, manifest: Manifest) -> Result<bool, Error> {
309        manifest.check_fresh(self.now(), team.signed())?;
310        let current = team.manifest();
311        if current.team.version > 0 {
312            manifest.may_replace(&current)?;
313        }
314        if manifest == *current {
315            return Ok(false);
316        }
317        tracing::info!(team = %team.config.id, version = manifest.team.version, "team manifest installed");
318        *team
319            .manifest
320            .write()
321            .unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::new(manifest);
322        Ok(true)
323    }
324
325    /// Reads a team's manifest again (from its file or URL).
326    pub(crate) async fn refresh(&self, team: &Team) -> Result<bool, Error> {
327        let manifest = match &team.config.source {
328            ManifestSource::Builtin => return Ok(false),
329            ManifestSource::File(path) => load_file(path, &team.config.id)?,
330            ManifestSource::Url { url, team_key } => {
331                let get = |u: String| async move {
332                    let reply = self
333                        .http
334                        .get(&u)
335                        .send()
336                        .await
337                        .and_then(reqwest::Response::error_for_status)
338                        .map_err(|e| Error::Manifest(format!("cannot fetch {u}: {e}")))?;
339                    reply
340                        .text()
341                        .await
342                        .map_err(|e| Error::Manifest(format!("cannot read {u}: {e}")))
343                };
344                let text = get(url.clone()).await?;
345                let sig = get(format!("{url}.sig")).await?;
346                let manifest = manifest::verify(&text, &sig, team_key)?;
347                if manifest.team.id != team.config.id {
348                    return Err(Error::Manifest(format!(
349                        "{url} is the manifest of team `{}`",
350                        manifest.team.id
351                    )));
352                }
353                manifest
354            }
355        };
356        self.install(team, manifest)
357    }
358
359    /// The public key of a member.
360    pub(crate) fn key_of(manifest: &Manifest, name: &str) -> Result<PublicKey, Error> {
361        manifest
362            .members
363            .get(name)
364            .ok_or_else(|| Error::Token(format!("`{name}` is not a member")))
365            .and_then(member_key)
366    }
367
368    /// The A2A version `member` of `team` speaks: 1.0 for okbots, as configured for agents.
369    pub(crate) fn version_of(&self, team: &Team, member: &str) -> a2a::Version {
370        if team.config.source == ManifestSource::Builtin {
371            return self
372                .external
373                .as_ref()
374                .and_then(|e| e.agent(member))
375                .map_or(a2a::Version::V10, |a| a.version);
376        }
377        a2a::Version::V10
378    }
379
380    pub(crate) fn is_down(&self, team: &str, member: &str) -> bool {
381        lock(&self.down)
382            .get(&(team.to_string(), member.to_string()))
383            .is_some_and(|until| *until > Instant::now())
384    }
385
386    fn mark(&self, team: &str, member: &str, up: bool) {
387        let key = (team.to_string(), member.to_string());
388        let mut down = lock(&self.down);
389        if up {
390            down.remove(&key);
391        } else {
392            down.insert(key, Instant::now() + DOWN_FOR);
393        }
394    }
395
396    /// Sends one JSON-RPC request to a member and returns its verified response.
397    pub(crate) async fn call(
398        &self,
399        team: &Team,
400        member: &str,
401        request: &a2a::Request,
402    ) -> Result<Value, PeerError> {
403        let result = self.exchange(team, member, request).await;
404        self.mark(
405            &team.config.id,
406            member,
407            !matches!(result, Err(PeerError::Unreachable(_))),
408        );
409        let response = result?;
410        if let Some(error) = response.get("error") {
411            let message = error
412                .get("message")
413                .and_then(Value::as_str)
414                .unwrap_or("refused")
415                .to_string();
416            return Err(PeerError::Refused(message));
417        }
418        response
419            .get("result")
420            .cloned()
421            .ok_or_else(|| PeerError::Refused("an empty response".into()))
422    }
423
424    async fn exchange(
425        &self,
426        team: &Team,
427        member: &str,
428        request: &a2a::Request,
429    ) -> Result<Value, PeerError> {
430        if team.config.source == ManifestSource::Builtin {
431            let agent = self
432                .external
433                .as_ref()
434                .and_then(|e| e.agent(member))
435                .ok_or_else(|| PeerError::NoRoute(member.to_string()))?;
436            return crate::external::call(&self.http, agent, request).await;
437        }
438        let manifest = team.manifest();
439        let target = manifest
440            .members
441            .get(member)
442            .ok_or_else(|| PeerError::NoRoute(member.to_string()))?;
443        let callee_key = member_key(target).map_err(|e| PeerError::Unreachable(e.to_string()))?;
444        let body = serde_json::to_vec(request).expect("requests serialize");
445        let claims = Claims::new(
446            &team.config.id,
447            &team.config.member,
448            member,
449            manifest.team.version,
450            &body,
451            self.now(),
452        );
453        let authorization = format!("{SCHEME}{}", token::sign(&self.key, &claims));
454        let (status, signature, reply) = match (&target.dir, &target.url, &target.relay) {
455            (Some(dir), _, _) => {
456                let socket = socket_path(&self.sockets, dir);
457                local_request(&socket, &team.config.id, &authorization, &body).await?
458            }
459            (_, Some(url), _) => self.http_request(url, &authorization, body).await?,
460            (_, _, Some(relay)) => {
461                Box::pin(crate::relay::send(
462                    self,
463                    team,
464                    relay,
465                    member,
466                    &authorization,
467                    &body,
468                ))
469                .await?
470            }
471            _ => return Err(PeerError::NoRoute(member.to_string())),
472        };
473        let Some(signature) = signature else {
474            return Err(PeerError::Refused(format!(
475                "`{member}` answered HTTP {status} without a signature"
476            )));
477        };
478        let claims = token::verify(&signature, &callee_key, &reply, self.now())
479            .map_err(|e| PeerError::Refused(format!("`{member}`'s answer: {e}")))?;
480        if claims.iss != member || claims.aud != team.config.member || claims.team != team.config.id
481        {
482            return Err(PeerError::Refused(format!(
483                "`{member}`'s answer is signed for someone else"
484            )));
485        }
486        serde_json::from_slice(&reply)
487            .map_err(|_| PeerError::Refused(format!("`{member}` sent invalid JSON")))
488    }
489
490    async fn http_request(
491        &self,
492        url: &str,
493        authorization: &str,
494        body: Vec<u8>,
495    ) -> Result<(u16, Option<String>, Vec<u8>), PeerError> {
496        let reply = self
497            .http
498            .post(url)
499            .header("authorization", authorization)
500            .header("content-type", "application/json")
501            .header(
502                crate::a2a::VERSION_HEADER,
503                crate::a2a::Version::V10.header(),
504            )
505            .header(crate::a2a::EXTENSIONS_HEADER, crate::a2a::EXTENSION)
506            .body(body)
507            .send()
508            .await
509            .map_err(|e| PeerError::Unreachable(without_url(e)))?;
510        let status = reply.status().as_u16();
511        let signature = reply
512            .headers()
513            .get(SIGNATURE_HEADER)
514            .and_then(|v| v.to_str().ok())
515            .map(str::to_string);
516        let bytes = reply
517            .bytes()
518            .await
519            .map_err(|e| PeerError::Unreachable(without_url(e)))?;
520        if bytes.len() > MAX_BODY {
521            return Err(PeerError::Refused("the answer is too large".into()));
522        }
523        Ok((status, signature, bytes.to_vec()))
524    }
525
526    /// Signs a response body for `caller`.
527    pub(crate) fn sign_response(&self, team: &Team, caller: &str, body: &[u8]) -> String {
528        let claims = Claims::new(
529            &team.config.id,
530            &team.config.member,
531            caller,
532            team.manifest().team.version,
533            body,
534            self.now(),
535        );
536        token::sign(&self.key, &claims)
537    }
538}
539
540fn without_url(error: reqwest::Error) -> String {
541    if error.is_timeout() {
542        return "timed out".into();
543    }
544    error.without_url().to_string()
545}
546
547/// A member's X25519 key: from the manifest, or from its bot directory.
548pub(crate) fn member_box_key(member: &Member) -> Result<crate::seal::BoxPublic, Error> {
549    match (&member.box_key, &member.dir) {
550        (Some(key), _) => Ok(*key),
551        (None, Some(dir)) => crate::seal::bot_box_key(dir),
552        (None, None) => Err(Error::Key(
553            "the member has no `box_key` (`okbot peer id` prints it)".into(),
554        )),
555    }
556}
557
558/// A member's key: from the manifest, or from its bot directory.
559pub(crate) fn member_key(member: &Member) -> Result<PublicKey, Error> {
560    match (&member.key, &member.dir) {
561        (Some(key), _) => Ok(key.clone()),
562        (None, Some(dir)) => bot_public_key(dir),
563        (None, None) => Err(Error::Key("the member has no key".into())),
564    }
565}
566
567fn load_file(path: &Path, team: &str) -> Result<Manifest, Error> {
568    let text =
569        std::fs::read_to_string(path).map_err(|e| Error::Io(path.display().to_string(), e))?;
570    let manifest = Manifest::parse(&text)?;
571    if manifest.team.id != team {
572        return Err(Error::Manifest(format!(
573            "{} is the manifest of team `{}`, not `{team}`",
574            path.display(),
575            manifest.team.id
576        )));
577    }
578    Ok(manifest)
579}
580
581/// One request over a Unix socket: a JSON line each way.
582#[derive(Debug, Serialize, Deserialize)]
583pub(crate) struct LocalRequest {
584    pub(crate) team: String,
585    pub(crate) authorization: String,
586    pub(crate) body: String,
587}
588
589/// The answer to a [`LocalRequest`].
590#[derive(Debug, Serialize, Deserialize)]
591pub(crate) struct LocalResponse {
592    pub(crate) status: u16,
593    #[serde(default)]
594    pub(crate) signature: Option<String>,
595    pub(crate) body: String,
596}
597
598async fn local_request(
599    socket: &Path,
600    team: &str,
601    authorization: &str,
602    body: &[u8],
603) -> Result<(u16, Option<String>, Vec<u8>), PeerError> {
604    let unreachable = |e: std::io::Error| {
605        PeerError::Unreachable(format!("{}: {e} (is that bot running?)", socket.display()))
606    };
607    let exchange = async {
608        let mut stream = tokio::net::UnixStream::connect(socket)
609            .await
610            .map_err(unreachable)?;
611        let request = LocalRequest {
612            team: team.to_string(),
613            authorization: authorization.to_string(),
614            body: String::from_utf8_lossy(body).into_owned(),
615        };
616        let mut line = serde_json::to_vec(&request).expect("requests serialize");
617        line.push(b'\n');
618        stream.write_all(&line).await.map_err(unreachable)?;
619        let mut reader = BufReader::new(stream).take(MAX_BODY as u64 * 2);
620        let mut answer = String::new();
621        reader.read_line(&mut answer).await.map_err(unreachable)?;
622        let response: LocalResponse = serde_json::from_str(&answer)
623            .map_err(|_| PeerError::Refused("an invalid local answer".into()))?;
624        Ok((
625            response.status,
626            response.signature,
627            response.body.into_bytes(),
628        ))
629    };
630    tokio::time::timeout(REQUEST_TIMEOUT, exchange)
631        .await
632        .map_err(|_| PeerError::Unreachable(format!("{}: timed out", socket.display())))?
633}
okbot_peer/node.rs

okbot_peer/
node.rs
