Publish harness and TUI open-source

initial sync from the monorepo

crates/common/xai-computer-hub-sdk/Cargo.toml

@@ -0,0 +1,57 @@
+[package]
+license = "Apache-2.0"
+name = "xai-computer-hub-sdk"
+version = "0.1.0"
+edition.workspace = true
+description = "SDK for the xAI Computer Hub: connection pool, transparent reconnect, tool harness, and tool-server runtime."
+
+[features]
+metrics = ["dep:prometheus"]
+
+[dependencies]
+tokio = { workspace = true, features = ["rt", "sync", "time", "macros"] }
+tokio-tungstenite = { workspace = true, features = ["rustls-tls-native-roots"] }
+tokio-util = { workspace = true }
+futures = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
+dashmap = { workspace = true }
+indexmap = { workspace = true }
+arc-swap = { workspace = true }
+async-trait = { workspace = true }
+thiserror = { workspace = true }
+tracing = { workspace = true }
+tracing-subscriber = { workspace = true }
+url = { workspace = true }
+http = { workspace = true }
+prometheus = { workspace = true, optional = true }
+reqwest = { workspace = true }
+chrono = { workspace = true }
+parking_lot = { workspace = true }
+
+fastrace = { workspace = true }
+# Trace donation: spans convert via the stock fastrace -> OTel reporter
+# and ship as standard OTLP payloads.
+fastrace-opentelemetry = { workspace = true }
+opentelemetry = { workspace = true }
+opentelemetry_sdk = { workspace = true }
+opentelemetry-proto = { workspace = true }
+prost = { workspace = true }
+base64 = { workspace = true }
+
+xai-tool-protocol = { workspace = true }
+xai-tool-runtime = { workspace = true }
+xai-tool-types = { workspace = true }
+xai-computer-hub-core = { workspace = true }
+xai-tracing = { workspace = true }
+
+# Integration tests that need heavier backend deps live in a separate sibling crate to keep this dev-dep set minimal.
+[dev-dependencies]
+tokio = { workspace = true, features = ["full", "test-util"] }
+axum = { workspace = true, features = ["ws", "macros"] }
+chrono = { workspace = true }
+base64 = { workspace = true }
+schemars = { workspace = true }
+
+[lints]
+workspace = true

crates/common/xai-computer-hub-sdk/src/admission.rs

@@ -0,0 +1,333 @@
+//! Three-tier semaphore admission + bounded-wait backpressure.
+//!
+//! Concurrent *running* calls are bounded at three scopes, acquired in a
+//! fixed **session → connection → global** order. A consistent
+//! most-local-first order is deadlock-free and never holds a scarce
+//! global permit while blocking on a local one. A single shared deadline
+//! spans all three acquisitions, so total admission latency is bounded by
+//! `wait_timeout`, not `3 × wait_timeout`.
+//!
+//! Under moderate pressure `admit` waits; under very high pressure the
+//! deadline elapses and the caller emits the shared overloaded JSON-RPC
+//! error (`-32016` "tool_busy") instead of silently dropping the request.
+
+use std::sync::{Arc, OnceLock};
+use std::time::Duration;
+
+use dashmap::DashMap;
+use serde_json::Value;
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use tokio::time::Instant;
+use xai_tool_protocol::{
+    JsonRpcError, JsonRpcId, JsonRpcResponse, JsonRpcVersion, ResponseOutcome, SessionId,
+};
+
+/// Numeric JSON-RPC code for overload rejection (`xai-tool-protocol`
+/// `error_codes.rs`: `-32016` "tool_busy").
+pub(crate) const TOOL_BUSY_CODE: i32 = -32016;
+
+const TOOL_BUSY_MESSAGE: &str = "tool server busy; tool call rejected";
+
+/// Default ceiling for the process-wide concurrency guard.
+pub(crate) const DEFAULT_GLOBAL_MAX_INFLIGHT: usize = 1024;
+/// Default per-session concurrent running calls.
+pub(crate) const DEFAULT_SESSION_MAX_INFLIGHT: usize = 16;
+/// Default per-connection concurrent running calls.
+pub(crate) const DEFAULT_CONN_MAX_INFLIGHT: usize = 256;
+/// Default bounded wait before an overloaded rejection.
+pub(crate) const DEFAULT_ADMISSION_WAIT_TIMEOUT: Duration = Duration::from_secs(3);
+
+/// Ops-tunable override for the process-wide global cap (Helm `env:`).
+const GLOBAL_MAX_INFLIGHT_ENV: &str = "XAI_TOOL_SERVER_GLOBAL_MAX_INFLIGHT";
+
+/// Inflight-gauge scope labels, in acquisition order. A held [`AdmitGuard`]
+/// counts against all three.
+const SCOPES: [&str; 3] = ["session", "conn", "global"];
+
+/// Build the shared overloaded (`-32016` "tool_busy") JSON-RPC error
+/// response. This is the single source of the overload wire shape, reused
+/// by BOTH the admission-timeout path (`server::execute_call`) and the
+/// demux inbox-full path (`demux::route_session`) so the two never drift.
+pub(crate) fn overloaded_response(id: JsonRpcId, session_id: SessionId) -> JsonRpcResponse<Value> {
+    JsonRpcResponse {
+        jsonrpc: JsonRpcVersion,
+        id,
+        session_id: Some(session_id),
+        outcome: ResponseOutcome::Error(JsonRpcError {
+            code: TOOL_BUSY_CODE,
+            message: TOOL_BUSY_MESSAGE.to_owned(),
+            data: Some(serde_json::json!({ "code": "tool_busy", "retryable": true })),
+        }),
+    }
+}
+
+/// Process-wide global admission semaphore, shared by every connection.
+///
+/// Initialized once at first use: the value comes from
+/// `XAI_TOOL_SERVER_GLOBAL_MAX_INFLIGHT` when present and parseable as a
+/// positive integer, otherwise `default_cap` (the builder knob, default
+/// [`DEFAULT_GLOBAL_MAX_INFLIGHT`]). Because the cell initializes exactly
+/// once, the first caller's `default_cap` and the env var at that instant
+/// fix the process-wide capacity.
+pub(crate) fn global_semaphore(default_cap: usize) -> Arc<Semaphore> {
+    static SEM: OnceLock<Arc<Semaphore>> = OnceLock::new();
+    SEM.get_or_init(|| {
+        let raw = std::env::var(GLOBAL_MAX_INFLIGHT_ENV).ok();
+        Arc::new(Semaphore::new(resolve_global_cap(
+            raw.as_deref(),
+            default_cap,
+        )))
+    })
+    .clone()
+}
+
+/// Resolve the process-wide global cap from the raw env value, falling
+/// back to `default_cap`. Pure (no global state) so the
+/// fall-back-never-panic guarantee is unit-tested: a non-numeric,
+/// negative, empty, or zero value all yield `default_cap`.
+fn resolve_global_cap(raw: Option<&str>, default_cap: usize) -> usize {
+    raw.and_then(|v| v.parse::<usize>().ok())
+        .filter(|n| *n > 0)
+        .unwrap_or(default_cap)
+}
+
+/// Why admission was refused.
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) enum Overloaded {
+    /// The bounded admission deadline elapsed under very high pressure.
+    Timeout,
+    /// A semaphore was closed — the server is shutting down.
+    Shutdown,
+}
+
+/// RAII guard holding all three permits for the call's lifetime.
+///
+/// Fields drop in declaration order, so permits are released in reverse
+/// of acquisition: global → connection → session.
+#[derive(Debug)]
+pub(crate) struct AdmitGuard {
+    _global: OwnedSemaphorePermit,
+    _conn: OwnedSemaphorePermit,
+    _session: OwnedSemaphorePermit,
+}
+
+impl Drop for AdmitGuard {
+    fn drop(&mut self) {
+        for scope in SCOPES {
+            crate::metrics::tool_call_inflight_dec(scope);
+        }
+    }
+}
+
+/// Three-tier admission controller. One per connection (`conn_sem`); the
+/// per-session map is created/destroyed alongside each session loop.
+#[derive(Debug)]
+pub(crate) struct Admission {
+    session_sems: DashMap<SessionId, Arc<Semaphore>>,
+    session_max: usize,
+    conn_sem: Arc<Semaphore>,
+    global_sem: Arc<Semaphore>,
+    wait_timeout: Duration,
+}
+
+impl Admission {
+    pub(crate) fn new(
+        session_max: usize,
+        conn_max: usize,
+        global_sem: Arc<Semaphore>,
+        wait_timeout: Duration,
+    ) -> Self {
+        Self {
+            session_sems: DashMap::new(),
+            session_max,
+            conn_sem: Arc::new(Semaphore::new(conn_max)),
+            global_sem,
+            wait_timeout,
+        }
+    }
+
+    /// Create the per-session semaphore entry. Called from
+    /// `bind_session_local` so the entry's lifetime is tied to the
+    /// session-loop task, not lazily minted in [`Self::admit`].
+    pub(crate) fn ensure_session(&self, session_id: &SessionId) {
+        self.session_sems
+            .entry(session_id.clone())
+            .or_insert_with(|| Arc::new(Semaphore::new(self.session_max)));
+    }
+
+    /// Remove the per-session semaphore entry on unbind / loop exit.
+    pub(crate) fn remove_session(&self, session_id: &SessionId) {
+        self.session_sems.remove(session_id);
+    }
+
+    /// Acquire one permit at each scope (session → connection → global)
+    /// against a single shared deadline.
+    pub(crate) async fn admit(&self, session_id: &SessionId) -> Result<AdmitGuard, Overloaded> {
+        let start = Instant::now();
+        let deadline = start + self.wait_timeout;
+        // The entry is created in `bind_session_local`; a straggler call
+        // admitted just after unbind cleanup falls back to a private,
+        // un-tracked semaphore rather than recreating a leaked entry.
+        let session_sem = self
+            .session_sems
+            .get(session_id)
+            .map(|s| s.clone())
+            .unwrap_or_else(|| Arc::new(Semaphore::new(self.session_max)));
+
+        let session = acquire_until(&session_sem, deadline).await?;
+        let conn = acquire_until(&self.conn_sem, deadline).await?;
+        let global = acquire_until(&self.global_sem, deadline).await?;
+
+        crate::metrics::admission_wait_observe(start.elapsed().as_secs_f64());
+        for scope in SCOPES {
+            crate::metrics::tool_call_inflight_inc(scope);
+        }
+        Ok(AdmitGuard {
+            _global: global,
+            _conn: conn,
+            _session: session,
+        })
+    }
+}
+
+/// Acquire one owned permit before `deadline`, mapping closed/elapsed to
+/// the matching [`Overloaded`] variant.
+async fn acquire_until(
+    sem: &Arc<Semaphore>,
+    deadline: Instant,
+) -> Result<OwnedSemaphorePermit, Overloaded> {
+    match tokio::time::timeout_at(deadline, sem.clone().acquire_owned()).await {
+        Ok(Ok(permit)) => Ok(permit),
+        Ok(Err(_closed)) => Err(Overloaded::Shutdown),
+        Err(_elapsed) => Err(Overloaded::Timeout),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn sid(s: &str) -> SessionId {
+        SessionId::new(s).expect("valid session id")
+    }
+
+    fn test_admission(session_max: usize, conn_max: usize, global_max: usize) -> Admission {
+        Admission::new(
+            session_max,
+            conn_max,
+            Arc::new(Semaphore::new(global_max)),
+            Duration::from_millis(150),
+        )
+    }
+
+    #[test]
+    fn resolve_global_cap_falls_back_on_bad_input_and_honors_valid() {
+        // Absent / non-numeric / negative / empty / zero → default (never panic).
+        assert_eq!(resolve_global_cap(None, 1024), 1024);
+        assert_eq!(resolve_global_cap(Some("abc"), 1024), 1024);
+        assert_eq!(resolve_global_cap(Some("-5"), 1024), 1024);
+        assert_eq!(resolve_global_cap(Some(""), 1024), 1024);
+        assert_eq!(resolve_global_cap(Some("0"), 1024), 1024);
+        assert_eq!(resolve_global_cap(Some("  7"), 1024), 1024); // leading space → parse fails
+        // A valid positive integer overrides the default.
+        assert_eq!(resolve_global_cap(Some("2048"), 1024), 2048);
+        assert_eq!(resolve_global_cap(Some("1"), 1024), 1);
+    }
+
+    #[test]
+    fn overloaded_response_carries_minus_32016_and_data_marker() {
+        let id: JsonRpcId = serde_json::from_value(serde_json::json!("call-1")).expect("id");
+        let resp = overloaded_response(id, sid("s1"));
+        let wire: Value = serde_json::from_str(&serde_json::to_string(&resp).expect("ser"))
+            .expect("round-trips to json");
+        assert_eq!(wire["error"]["code"], TOOL_BUSY_CODE);
+        assert_eq!(wire["error"]["code"], -32016);
+        assert_eq!(wire["error"]["data"]["code"], "tool_busy");
+        assert_eq!(wire["error"]["data"]["retryable"], true);
+        assert_eq!(wire["session_id"], "s1");
+        assert_eq!(wire["id"], "call-1");
+        assert!(
+            wire.get("result").is_none(),
+            "overload is an error, never a result"
+        );
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn admit_times_out_when_session_saturated() {
+        let admission = test_admission(2, 16, 64);
+        let session = sid("sat");
+        admission.ensure_session(&session);
+
+        // Hold both session permits.
+        let g1 = admission.admit(&session).await.expect("first admit");
+        let _g2 = admission.admit(&session).await.expect("second admit");
+
+        // Third admit must elapse the deadline → Timeout (not a hang).
+        let result = admission.admit(&session).await;
+        assert_eq!(result.unwrap_err(), Overloaded::Timeout);
+
+        // Releasing one permit frees a slot for the next admit.
+        drop(g1);
+        admission
+            .admit(&session)
+            .await
+            .expect("permit released → admit succeeds");
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn admit_blocks_on_connection_scope_when_conn_saturated() {
+        // conn_max = 1 is the binding constraint even though session has
+        // room; a second admit on a *different* session still times out.
+        let admission = test_admission(8, 1, 64);
+        let a = sid("a");
+        let b = sid("b");
+        admission.ensure_session(&a);
+        admission.ensure_session(&b);
+
+        let _held = admission.admit(&a).await.expect("first admit");
+        let result = admission.admit(&b).await;
+        assert_eq!(
+            result.unwrap_err(),
+            Overloaded::Timeout,
+            "connection cap binds across sessions"
+        );
+    }
+
+    #[tokio::test]
+    async fn admit_succeeds_repeatedly_under_capacity() {
+        let admission = test_admission(4, 16, 64);
+        let session = sid("ok");
+        admission.ensure_session(&session);
+        let mut guards = Vec::new();
+        for _ in 0..4 {
+            guards.push(admission.admit(&session).await.expect("within capacity"));
+        }
+        assert_eq!(guards.len(), 4);
+    }
+
+    #[tokio::test]
+    async fn closed_semaphore_maps_to_shutdown() {
+        let global = Arc::new(Semaphore::new(0));
+        global.close();
+        let admission = Admission::new(4, 16, global, Duration::from_secs(5));
+        let session = sid("closed");
+        admission.ensure_session(&session);
+        let result = admission.admit(&session).await;
+        assert_eq!(result.unwrap_err(), Overloaded::Shutdown);
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn straggler_admit_after_remove_uses_private_permit() {
+        let admission = test_admission(1, 16, 64);
+        let session = sid("gone");
+        // No ensure_session: simulate a straggler after unbind removed it.
+        admission.remove_session(&session);
+        // Falls back to a private semaphore and still admits (no panic,
+        // no leaked tracked entry).
+        let _g = admission.admit(&session).await.expect("private fallback");
+        assert!(
+            admission.session_sems.get(&session).is_none(),
+            "straggler must not recreate a tracked entry"
+        );
+    }
+}

crates/common/xai-computer-hub-sdk/src/auth.rs

@@ -0,0 +1,238 @@
+//! Auth credentials and pool-dedup principal keys.
+//!
+//! [`AuthCredential`] models the credential the client attaches at
+//! handshake time. Two variants are supported:
+//!
+//! - [`AuthCredential::Bearer`] for the simple "Authorization: Bearer
+//!   …" path (e.g. JWT-against-OAuth2 deployments).
+//! - [`AuthCredential::Headers`] for callers that already hold a
+//!   pre-built header bundle (e.g. signed identity headers generated
+//!   by an upstream proxy or test harness).
+//!
+//! [`PrincipalKey`] is the stable hashable projection of an
+//! `AuthCredential`; the pool keys connections by
+//! `(url, principal_key)` so two [`crate::ToolServer`] builds with the
+//! same credential reuse one socket while distinct credentials open
+//! distinct sockets. The server derives `user_id` from the credential at
+//! upgrade time and returns it in the hello ack — the SDK never needs
+//! to carry `user_id` alongside the credential.
+//!
+//! ## Pool dedup and credential refresh
+//!
+//! Both variants include the secret material in the `PrincipalKey`
+//! fingerprint. This is deliberate: distinct secrets imply distinct
+//! credentials, so two callers with different tokens open distinct
+//! sockets. The trade-off is that a caller that rotates its bearer JWT
+//! every N minutes will open a new socket on each rotation.
+//! Long-running tool servers should reuse the SAME [`AuthCredential`]
+//! instance across builds and refresh the credential out-of-band rather
+//! than hand a fresh JWT to every build.
+
+use std::collections::BTreeMap;
+use std::fmt;
+
+use http::HeaderName;
+use http::header::AUTHORIZATION;
+
+use crate::error::ClientError;
+
+/// Credential carried into the WebSocket upgrade.
+///
+/// Clones are cheap (the secret material is at most a small number of
+/// owned strings). The server derives `user_id` from the credential at
+/// upgrade time and returns it in the [`xai_tool_protocol::HelloAckMsg`].
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub enum AuthCredential {
+    /// Bearer token attached as the `Authorization: Bearer …` header.
+    Bearer { token: String },
+    /// Pre-built header bundle. Used when the auth flow lives outside
+    /// the SDK (e.g. an upstream proxy that already produced signed
+    /// identity headers). Header order is canonicalised for stable
+    /// hashing via [`BTreeMap`]; names are lowercased and validated
+    /// as `HeaderName` at construction time so an invalid name
+    /// surfaces as [`ClientError::InvalidConfig`] instead of being
+    /// silently dropped at upgrade time.
+    Headers { headers: BTreeMap<String, String> },
+}
+
+impl AuthCredential {
+    /// Convenience constructor for the bearer-token shape.
+    pub fn bearer(token: impl Into<String>) -> Self {
+        Self::Bearer {
+            token: token.into(),
+        }
+    }
+
+    /// Convenience constructor for the raw-header bundle shape.
+    ///
+    /// Names are canonicalised to lowercase and validated as
+    /// [`HeaderName`] at construction so an invalid header (e.g. one
+    /// containing a newline injection attempt) returns
+    /// [`ClientError::InvalidConfig`] rather than being silently
+    /// filtered out at upgrade time.
+    pub fn headers<I, K, V>(headers: I) -> Result<Self, ClientError>
+    where
+        I: IntoIterator<Item = (K, V)>,
+        K: AsRef<str>,
+        V: Into<String>,
+    {
+        let mut map: BTreeMap<String, String> = BTreeMap::new();
+        for (raw_name, raw_value) in headers {
+            let name = raw_name.as_ref().to_ascii_lowercase();
+            HeaderName::from_bytes(name.as_bytes()).map_err(|err| {
+                ClientError::InvalidConfig(format!("invalid header name {name:?}: {err}"))
+            })?;
+            map.insert(name, raw_value.into());
+        }
+        Ok(Self::Headers { headers: map })
+    }
+
+    /// Stable hashable projection used as the pool dedup key.
+    ///
+    /// Distinct credentials hash equal iff they carry the same secret
+    /// material. See the module-level "Pool dedup and credential
+    /// refresh" section for the implications when bearer tokens are
+    /// rotated.
+    pub fn principal_key(&self) -> PrincipalKey {
+        match self {
+            Self::Bearer { token } => PrincipalKey {
+                fingerprint: format!("bearer:{token}"),
+            },
+            Self::Headers { headers } => {
+                // Concatenate canonicalised name=value pairs so the
+                // fingerprint is order-independent.
+                let mut joined = String::with_capacity(headers.len() * 32);
+                for (name, value) in headers {
+                    joined.push_str(name);
+                    joined.push('=');
+                    joined.push_str(value);
+                    joined.push('\n');
+                }
+                PrincipalKey {
+                    fingerprint: format!("headers:{joined}"),
+                }
+            }
+        }
+    }
+
+    /// Headers to attach to the WebSocket upgrade request.
+    ///
+    /// `Headers` variant entries are infallible at this point — names
+    /// were validated by [`Self::headers`].
+    pub fn upgrade_headers(&self) -> Vec<(HeaderName, String)> {
+        match self {
+            Self::Bearer { token, .. } => {
+                vec![(AUTHORIZATION, format!("Bearer {token}"))]
+            }
+            Self::Headers { headers, .. } => headers
+                .iter()
+                .filter_map(|(name, value)| {
+                    HeaderName::from_bytes(name.as_bytes())
+                        .ok()
+                        .map(|n| (n, value.clone()))
+                })
+                .collect(),
+        }
+    }
+}
+
+impl fmt::Debug for AuthCredential {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // Never log the secret; surface only the variant.
+        match self {
+            Self::Bearer { .. } => f
+                .debug_struct("AuthCredential::Bearer")
+                .finish_non_exhaustive(),
+            Self::Headers { headers } => f
+                .debug_struct("AuthCredential::Headers")
+                .field("header_count", &headers.len())
+                .finish_non_exhaustive(),
+        }
+    }
+}
+
+/// Stable hashable projection of an [`AuthCredential`] used as the
+/// pool dedup key alongside the connect URL. Two connections with the
+/// same token fingerprint will get the same server-assigned `user_id`.
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct PrincipalKey {
+    fingerprint: String,
+}
+
+impl fmt::Debug for PrincipalKey {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("PrincipalKey").finish_non_exhaustive()
+    }
+}
+
+/// Owner identity surfaced by an [`AuthProvider`] alongside its credential.
+///
+/// Mirrors the OAuth principal fields the provider parsed from its auth source.
+/// It is kept separate from [`AuthCredential`] on purpose: identity must NOT
+/// participate in pool-dedup hashing (that keys only on the secret), and the
+/// credential's `Eq`/`Hash` derives must stay token-only. Consumers (e.g. the
+/// workspace) map this onto their own identity record.
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct AuthIdentity {
+    /// Stable user identifier (owner of the bearer token).
+    pub user_id: String,
+    /// OAuth `principal_type` wire string (`"User"` / `"Team"`), when known.
+    pub principal_type: Option<String>,
+    /// Team id when `principal_type == "Team"`; otherwise `None`.
+    pub principal_id: Option<String>,
+}
+
+/// Credential provider called on every connect/reconnect.
+pub trait AuthProvider: Send + Sync + std::fmt::Debug {
+    fn current(&self) -> AuthCredential;
+
+    /// Stable pool-dedup key, decoupled from the per-connect credential.
+    ///
+    /// Defaults to the current credential's key (existing behavior). A provider
+    /// that re-mints a rotating secret on every [`Self::current`] call (e.g. a
+    /// refresh-before-use bearer) MUST override this to key only on stable
+    /// identity, otherwise each rotation fragments the connection pool.
+    fn principal_key(&self) -> PrincipalKey {
+        self.current().principal_key()
+    }
+
+    /// Owner identity behind the credential, when the provider can surface it.
+    ///
+    /// Defaults to `None` for providers that only carry a bearer token (e.g. a
+    /// bare [`AuthCredential`]). Providers that parse OAuth principal fields
+    /// (e.g. OIDC) override this so downstream consumers can attribute
+    /// requests without a second auth-source read.
+    fn identity(&self) -> Option<AuthIdentity> {
+        None
+    }
+}
+
+pub type SharedAuthProvider = std::sync::Arc<dyn AuthProvider>;
+
+impl AuthProvider for AuthCredential {
+    fn current(&self) -> AuthCredential {
+        self.clone()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn invalid_header_name_rejected_at_construction() {
+        let cred = AuthCredential::headers([("authorization\nx-injected", "value")]);
+        match cred {
+            Err(ClientError::InvalidConfig(msg)) => {
+                assert!(msg.contains("invalid header name"), "got {msg}")
+            }
+            other => panic!("expected InvalidConfig; got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn valid_headers_accepted() {
+        let cred = AuthCredential::headers([("authorization", "Bearer token")]).expect("valid");
+        assert_eq!(cred.upgrade_headers().len(), 1);
+    }
+}

crates/common/xai-computer-hub-sdk/src/cancel.rs

@@ -0,0 +1,301 @@
+//! Per-session strict-cancellation registry.
+//!
+//! Maps each in-flight `tool_call_id` to its [`CancellationToken`] so a
+//! `Cancel` hook (or session teardown) can hard-cancel the running call
+//! by dropping its future. A small `pending` tombstone set covers the
+//! race where a `Cancel` arrives *before* the dispatcher registered the
+//! token (the symmetric window to pre-spawn registration): the id is
+//! tombstoned and the dispatcher cancels it at registration time.
+//!
+//! One registry per session, tied to the session-loop lifetime alongside
+//! the inbox and the per-session admission semaphore.
+
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use dashmap::{DashMap, DashSet};
+use tokio_util::sync::CancellationToken;
+use xai_tool_protocol::ToolCallId;
+
+/// Upper bound on outstanding pre-registration tombstones. Tombstones
+/// cover the microscopic window between a `Cancel` hook and the matching
+/// `register`, so in steady state the set holds a handful of entries. A
+/// `Cancel` whose call never registers (e.g. one racing call completion,
+/// after `deregister` already removed the live token) leaves a tombstone
+/// that no `register` ever consumes; this cap reclaims such stragglers so
+/// a single long-lived session cannot grow `pending` without bound.
+const MAX_PENDING_TOMBSTONES: usize = 8192;
+
+/// Per-session `tool_call_id -> CancellationToken` map plus a pending
+/// tombstone set for cancels that land before registration.
+#[derive(Default, Debug)]
+pub(crate) struct CancelRegistry {
+    map: DashMap<ToolCallId, CancellationToken>,
+    pending: DashSet<ToolCallId>,
+    /// Set once by [`Self::cancel_all`] (teardown). After this, every new
+    /// `register` starts cancelled so a request dispatched in the teardown
+    /// window cannot escape as an orphaned, uncancellable task.
+    closed: AtomicBool,
+}
+
+impl CancelRegistry {
+    /// Register `token` for `call_id` before the call is spawned. If a
+    /// `Cancel` already tombstoned this id, the token is cancelled
+    /// immediately so the call starts cancelled. Returns whether the
+    /// token was pre-cancelled (by a tombstone or because the registry was
+    /// torn down).
+    pub(crate) fn register(&self, call_id: ToolCallId, token: &CancellationToken) -> bool {
+        if self.closed.load(Ordering::Acquire) {
+            token.cancel();
+            return true;
+        }
+        let pre_cancelled = self.pending.remove(&call_id).is_some();
+        if pre_cancelled {
+            token.cancel();
+        }
+        self.map.insert(call_id.clone(), token.clone());
+        // Re-check after the insert: if `cancel_all` drained the map
+        // between our closed-check and the insert, our entry would be
+        // missed. The DashMap shard lock orders the insert against the
+        // drain, so observing `closed` here guarantees we cancel + drop
+        // any entry the drain could not reach (closes the teardown race).
+        if self.closed.load(Ordering::Acquire) {
+            if let Some((_, missed)) = self.map.remove(&call_id) {
+                missed.cancel();
+            }
+            return true;
+        }
+        pre_cancelled
+    }
+
+    /// Cancel a live call, else tombstone the id so the dispatcher
+    /// cancels it at registration time. Returns true when a live token
+    /// was found and cancelled.
+    pub(crate) fn cancel(&self, call_id: &ToolCallId) -> bool {
+        if let Some((_, token)) = self.map.remove(call_id) {
+            token.cancel();
+            true
+        } else {
+            if self.pending.len() >= MAX_PENDING_TOMBSTONES {
+                // Evict one straggler tombstone (a cancel whose call never
+                // registered) before inserting so the set stays bounded.
+                // Collect the key first, then remove, so we never hold a
+                // shard iterator across the removal.
+                let stale = self.pending.iter().next().map(|e| e.key().clone());
+                if let Some(stale) = stale {
+                    self.pending.remove(&stale);
+                }
+            }
+            self.pending.insert(call_id.clone());
+            false
+        }
+    }
+
+    /// Deregister a call's token on completion or cancel. Idempotent.
+    pub(crate) fn deregister(&self, call_id: &ToolCallId) {
+        self.map.remove(call_id);
+    }
+
+    /// Whether [`Self::cancel_all`] has closed this registry. A closed
+    /// registry marks a session whose loop is (or is about to be) torn
+    /// down — used by the soft-rebind liveness gate.
+    pub(crate) fn is_closed(&self) -> bool {
+        self.closed.load(Ordering::Acquire)
+    }
+
+    /// Drain-and-cancel every live token and close the registry. Used on
+    /// session teardown (`unbind_session` / `shutdown` / full rebind of a
+    /// dead loop — a soft rebind of a live session keeps its registry) so
+    /// detached `execute_call` tasks wind down promptly AND any call
+    /// dispatched in
+    /// the teardown window starts cancelled (see [`Self::register`]).
+    /// Returns the number of tokens cancelled.
+    pub(crate) fn cancel_all(&self) -> usize {
+        // Mark closed BEFORE draining so a concurrent `register` either
+        // observes the close (and self-cancels) or has its entry drained
+        // here — never both-miss.
+        self.closed.store(true, Ordering::Release);
+        let mut cancelled = 0;
+        self.map.retain(|_, token| {
+            token.cancel();
+            cancelled += 1;
+            false
+        });
+        // Drop tombstones too: teardown closes the registry, so no future
+        // `register` will consume them. Leaving them would let a stale
+        // straggler set survive to the end of the (already-done) session.
+        self.pending.clear();
+        cancelled
+    }
+
+    #[cfg(test)]
+    pub(crate) fn live_count(&self) -> usize {
+        self.map.len()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn pending_count(&self) -> usize {
+        self.pending.len()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn cid() -> ToolCallId {
+        ToolCallId::new_v7()
+    }
+
+    #[test]
+    fn cancel_live_token_fires_and_removes_entry() {
+        let reg = CancelRegistry::default();
+        let id = cid();
+        let token = CancellationToken::new();
+        assert!(
+            !reg.register(id.clone(), &token),
+            "fresh register, no tombstone"
+        );
+        assert_eq!(reg.live_count(), 1);
+
+        assert!(reg.cancel(&id), "live token must report a hit");
+        assert!(
+            token.is_cancelled(),
+            "the registered token must be cancelled"
+        );
+        assert_eq!(reg.live_count(), 0, "cancel removes the live entry");
+        assert_eq!(reg.pending_count(), 0, "a live hit leaves no tombstone");
+    }
+
+    #[test]
+    fn cancel_before_registration_tombstones_then_register_pre_cancels() {
+        let reg = CancelRegistry::default();
+        let id = cid();
+
+        // Cancel arrives first: no live token, so it tombstones.
+        assert!(!reg.cancel(&id), "no live token yet → miss");
+        assert_eq!(reg.pending_count(), 1);
+        assert_eq!(reg.live_count(), 0);
+
+        // Registration consumes the tombstone and starts cancelled.
+        let token = CancellationToken::new();
+        assert!(
+            reg.register(id.clone(), &token),
+            "register must report the pre-cancel"
+        );
+        assert!(token.is_cancelled(), "tombstone must pre-cancel the token");
+        assert_eq!(reg.pending_count(), 0, "tombstone consumed at registration");
+        assert_eq!(reg.live_count(), 1);
+    }
+
+    #[test]
+    fn deregister_clears_live_entry_without_cancel() {
+        let reg = CancelRegistry::default();
+        let id = cid();
+        let token = CancellationToken::new();
+        reg.register(id.clone(), &token);
+
+        reg.deregister(&id);
+        assert_eq!(reg.live_count(), 0);
+        assert!(
+            !token.is_cancelled(),
+            "deregister on normal completion must NOT cancel the token"
+        );
+        // A later cancel for a completed call only tombstones (harmless).
+        assert!(!reg.cancel(&id));
+        assert_eq!(reg.pending_count(), 1);
+    }
+
+    #[test]
+    fn cancel_all_drains_and_cancels_every_live_token() {
+        let reg = CancelRegistry::default();
+        let ids: Vec<ToolCallId> = (0..5).map(|_| cid()).collect();
+        let tokens: Vec<CancellationToken> = ids
+            .iter()
+            .map(|id| {
+                let t = CancellationToken::new();
+                reg.register(id.clone(), &t);
+                t
+            })
+            .collect();
+        assert_eq!(reg.live_count(), 5);
+
+        assert_eq!(
+            reg.cancel_all(),
+            5,
+            "cancel_all reports every drained token"
+        );
+        assert_eq!(reg.live_count(), 0, "registry is empty after teardown");
+        for token in &tokens {
+            assert!(token.is_cancelled(), "every live token must be cancelled");
+        }
+        // Idempotent: a second teardown cancels nothing.
+        assert_eq!(reg.cancel_all(), 0);
+    }
+
+    #[test]
+    fn register_after_cancel_all_starts_cancelled() {
+        // Teardown race regression: once `cancel_all` has closed the
+        // registry, a call dispatched in the teardown window must start
+        // cancelled and must NOT linger as a live, uncancellable entry.
+        let reg = CancelRegistry::default();
+        assert_eq!(reg.cancel_all(), 0, "empty teardown cancels nothing");
+
+        let id = cid();
+        let token = CancellationToken::new();
+        assert!(
+            reg.register(id.clone(), &token),
+            "register on a closed registry must report pre-cancel"
+        );
+        assert!(
+            token.is_cancelled(),
+            "a call dispatched after teardown must start cancelled"
+        );
+        assert_eq!(
+            reg.live_count(),
+            0,
+            "a closed-registry register must not leave a live (orphan) entry"
+        );
+    }
+
+    #[test]
+    fn register_without_tombstone_does_not_cancel() {
+        let reg = CancelRegistry::default();
+        let id = cid();
+        let token = CancellationToken::new();
+        assert!(!reg.register(id, &token));
+        assert!(
+            !token.is_cancelled(),
+            "a clean registration must leave the token live"
+        );
+    }
+
+    #[test]
+    fn cancel_all_clears_pending_tombstones() {
+        let reg = CancelRegistry::default();
+        reg.cancel(&cid());
+        reg.cancel(&cid());
+        assert_eq!(reg.pending_count(), 2);
+        reg.cancel_all();
+        assert_eq!(
+            reg.pending_count(),
+            0,
+            "teardown must drop pending tombstones"
+        );
+    }
+
+    #[test]
+    fn pending_tombstones_stay_bounded_under_spurious_cancels() {
+        // A long-lived session that keeps receiving cancels for call_ids
+        // that never register (e.g. cancels racing call completion) must
+        // not grow `pending` without bound.
+        let reg = CancelRegistry::default();
+        for _ in 0..(MAX_PENDING_TOMBSTONES + 256) {
+            assert!(!reg.cancel(&cid()), "never-registered id is a miss");
+        }
+        assert!(
+            reg.pending_count() <= MAX_PENDING_TOMBSTONES,
+            "tombstone set must stay within its cap, got {}",
+            reg.pending_count()
+        );
+    }
+}

crates/common/xai-computer-hub-sdk/src/connection_borrow.rs

@@ -0,0 +1,224 @@
+//! Shared connection-borrow lifecycle for `ToolServer` and `ToolHarness`.
+//!
+//! Wraps a pooled [`HubConnection`] with a [`CancellationToken`] for
+//! shutdown coordination and an at-most-once `torn_down` guard.
+
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use tokio_util::sync::CancellationToken;
+use url::Url;
+use xai_tool_protocol::ConnectionKind;
+
+use crate::auth::AuthProvider;
+use crate::connection::{
+    ConnectCallback, ConnectionTuning, DisconnectCallback, HubConnection, ReconnectCallback,
+};
+use crate::error::ClientError;
+use crate::pool::HubConnectionPool;
+
+/// Borrowed slice of a pooled [`HubConnection`] plus the refcount of
+/// session bindings the borrower owns. Drop guard lives here so the
+/// teardown sequence is at-most-once across explicit `shutdown` and
+/// the `Drop` fallback.
+pub(crate) struct ConnectionBorrow {
+    connection: Arc<HubConnection>,
+    shutdown: CancellationToken,
+    /// At-most-once guard coordinated via `compare_exchange`.
+    torn_down: AtomicBool,
+}
+
+impl std::fmt::Debug for ConnectionBorrow {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ConnectionBorrow")
+            .field(
+                "torn_down",
+                &self.torn_down.load(std::sync::atomic::Ordering::Relaxed),
+            )
+            .finish_non_exhaustive()
+    }
+}
+
+impl ConnectionBorrow {
+    /// Resolve a pool entry, refcount-bind every requested session,
+    /// and return a borrow. On any per-session bind failure the
+    /// already-bound sessions are unregistered before returning the
+    /// error so partial state never leaks.
+    pub(crate) async fn acquire(
+        pool: Arc<HubConnectionPool>,
+        url: Url,
+        auth: Arc<dyn AuthProvider>,
+        kind: ConnectionKind,
+        on_reconnect: Option<Arc<ReconnectCallback>>,
+        on_disconnect: Option<Arc<DisconnectCallback>>,
+        on_connect: Option<Arc<ConnectCallback>>,
+        server_id: Option<xai_tool_protocol::ServerId>,
+        server_description: Option<String>,
+        server_metadata: Option<serde_json::Value>,
+        alpha_test_key: Option<String>,
+        allow_insecure_ws: bool,
+        tuning: ConnectionTuning,
+    ) -> Result<Self, ClientError> {
+        let connection = pool
+            .get_or_connect_tuned(
+                url,
+                auth,
+                kind,
+                on_reconnect,
+                on_disconnect,
+                on_connect,
+                server_id,
+                server_description,
+                server_metadata,
+                alpha_test_key,
+                allow_insecure_ws,
+                tuning,
+            )
+            .await?;
+        Ok(Self {
+            connection,
+            shutdown: CancellationToken::new(),
+            torn_down: AtomicBool::new(false),
+        })
+    }
+
+    pub(crate) fn connection(&self) -> &Arc<HubConnection> {
+        &self.connection
+    }
+
+    pub(crate) fn shutdown_token(&self) -> &CancellationToken {
+        &self.shutdown
+    }
+
+    /// Returns `true` if this caller won the at-most-once teardown.
+    pub(crate) fn begin_teardown(&self) -> bool {
+        self.torn_down
+            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
+            .is_ok()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use std::net::SocketAddr;
+    use std::sync::Arc;
+
+    use crate::auth::AuthCredential;
+    use axum::Router;
+    use axum::extract::WebSocketUpgrade;
+    use axum::extract::ws::{Message, WebSocket};
+    use axum::response::IntoResponse;
+    use axum::routing::get;
+    use serde_json::json;
+    use tokio::net::TcpListener;
+
+    /// Spawn an in-process mock server that completes the WebSocket
+    /// handshake and ignores everything else. Returned address is
+    /// bound on `127.0.0.1`.
+    async fn spawn_borrow_mock_hub() -> SocketAddr {
+        let app = Router::new().route("/v1/tools", get(ws_upgrade));
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind ephemeral");
+        let addr = listener.local_addr().expect("local addr");
+        tokio::spawn(async move {
+            let _ = axum::serve(listener, app.into_make_service()).await;
+        });
+        tokio::task::yield_now().await;
+        addr
+    }
+
+    async fn ws_upgrade(ws: WebSocketUpgrade) -> impl IntoResponse {
+        ws.on_upgrade(handle_socket)
+    }
+
+    async fn handle_socket(mut socket: WebSocket) {
+        let _ = socket.recv().await;
+        let ack = json!({
+            "connection_id": "borrow-mock",
+            "user_id": "test",
+            "computer_hub_version": "test",
+            "supported_protocol_versions": ["1.0.0"],
+        });
+        let _ = socket.send(Message::Text(ack.to_string().into())).await;
+        // Keep the WebSocket alive until the client disconnects.
+        // These tests only exercise borrow lifecycle (teardown
+        // atomicity), not protocol frames.
+        while let Some(Ok(_msg)) = socket.recv().await {}
+    }
+
+    async fn acquire_borrow() -> ConnectionBorrow {
+        let addr = spawn_borrow_mock_hub().await;
+        let url = Url::parse(&format!("ws://{addr}/v1/tools")).expect("valid url");
+        let cred: Arc<dyn AuthProvider> = Arc::new(AuthCredential::bearer("ignored"));
+        let pool = HubConnectionPool::new();
+        ConnectionBorrow::acquire(
+            pool,
+            url,
+            cred,
+            ConnectionKind::Harness,
+            None, // on_reconnect
+            None, // on_disconnect
+            None, // on_connect
+            None, // server_id
+            None, // server_description
+            None, // server_metadata
+            None, // alpha_test_key
+            false,
+            ConnectionTuning::default(),
+        )
+        .await
+        .expect("acquire borrow")
+    }
+
+    #[tokio::test]
+    async fn begin_teardown_returns_true_once_and_false_after() {
+        let borrow = acquire_borrow().await;
+        assert!(
+            borrow.begin_teardown(),
+            "first call wins the at-most-once transition"
+        );
+        assert!(
+            !borrow.begin_teardown(),
+            "subsequent calls observe the already-torn-down state"
+        );
+        assert!(
+            !borrow.begin_teardown(),
+            "the at-most-once transition is sticky"
+        );
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+    async fn begin_teardown_is_atomic_under_concurrent_callers() {
+        let borrow = Arc::new(acquire_borrow().await);
+        let n_callers = 64;
+        let barrier = Arc::new(tokio::sync::Barrier::new(n_callers));
+        let mut handles = Vec::with_capacity(n_callers);
+        for _ in 0..n_callers {
+            let borrow = borrow.clone();
+            let barrier = barrier.clone();
+            handles.push(tokio::spawn(async move {
+                barrier.wait().await;
+                borrow.begin_teardown()
+            }));
+        }
+        let mut wins = 0usize;
+        for h in handles {
+            if h.await.expect("join") {
+                wins += 1;
+            }
+        }
+        assert_eq!(
+            wins, 1,
+            "exactly one of {n_callers} concurrent callers must win the at-most-once transition"
+        );
+    }
+
+    #[tokio::test]
+    async fn acquire_returns_zero_bound_sessions() {
+        let borrow = acquire_borrow().await;
+        assert_eq!(borrow.connection().bound_session_count(), 0);
+    }
+}

crates/common/xai-computer-hub-sdk/src/demux.rs

@@ -0,0 +1,973 @@
+//! Inbound frame demultiplexer.
+//!
+//! Frames inbound from the WebSocket fall into four buckets:
+//!
+//! 1. JSON-RPC **responses** correlated to a previously-issued request
+//!    by `id`. Routed through the crate-internal response-waiter map.
+//! 2. **`tool_call_progress` notifications** correlated to a per-call
+//!    `tool_call_id` carried in `params`. Routed through the
+//!    crate-internal progress-waiter map registered via
+//!    `Demux::try_register_progress_waiter` (crate-internal).
+//! 3. JSON-RPC **requests / notifications** carrying a `session_id` —
+//!    routed to the per-session inbox registered via
+//!    [`Demux::register_session_inbox`].
+//! 4. Connection-level frames (handshake, ping/pong) that the
+//!    connection actor handles directly without going through the demux.
+//!
+//! The demux owns the session inbox map, the in-flight response
+//! waiters, and the per-call progress waiters; the connection actor
+//! parses each text frame, classifies it,
+//! and pushes it through this module.
+//!
+//! Routing inbound frames is non-blocking: a full session inbox or a
+//! dropped receiver returns a typed [`RouteOutcome`] variant rather
+//! than awaiting the inbox. Blocking on a slow consumer would back up
+//! the entire connection actor and starve every other session sharing
+//! the socket.
+
+use dashmap::DashMap;
+use serde_json::Value;
+use tokio::sync::oneshot;
+use tracing::warn;
+
+use xai_tool_protocol::{
+    JsonRpcId, JsonRpcResponse, RequestId, SessionId, ToolCallId, ToolCallProgressFrame,
+};
+
+use crate::error::ClientError;
+
+/// Frame routed to a session inbox.
+#[derive(Debug, Clone)]
+pub enum InboundFrame {
+    /// A request the inbox owner must answer (any session frame carrying
+    /// an `id`): a `tool_call_request`, or a reverse-direction `hook`
+    /// answered via `ToolHarness::send_hook_reply`. Carries raw JSON.
+    Request(Value),
+    /// Server-issued notification (e.g. `tool.notification`) — fire-and-
+    /// forget, no reply expected.
+    Notification(Value),
+}
+
+/// Outcome of [`Demux::route`].
+#[derive(Debug, PartialEq, Eq)]
+pub enum RouteOutcome {
+    /// Matched a response waiter; the oneshot was fulfilled.
+    Response,
+    /// Forwarded to a session inbox.
+    Session,
+    /// Matched a progress waiter; the progress frame was forwarded to
+    /// the per-call progress channel.
+    Progress,
+    /// No inbox is bound for the targeted session.
+    UnknownSession,
+    /// No progress waiter is parked for the targeted `tool_call_id`. The
+    /// caller's stream is no longer subscribed (typical post-terminal),
+    /// so the frame is dropped.
+    UnknownProgress,
+    /// Connection-level notification broadcast to subscribers.
+    Notification,
+    /// No waiter is parked for the targeted request id, OR the frame
+    /// was unaddressable.
+    Unrouted,
+    /// The session inbox sender was full; the frame was dropped to
+    /// avoid blocking the connection actor.
+    InboxFull,
+    /// The session inbox receiver was dropped (e.g. the consumer's
+    /// run loop exited); the binding is now stale and the frame was
+    /// dropped.
+    SessionDropped,
+    /// The progress channel was full; the frame was dropped to avoid
+    /// blocking the connection actor. The caller's stream consumer
+    /// fell behind on draining progress.
+    ProgressFull,
+    /// The progress receiver was dropped (e.g. the caller's stream was
+    /// dropped); the waiter binding is now stale and the frame was
+    /// dropped.
+    ProgressDropped,
+}
+
+/// Demux state. Cheap to construct; uses [`DashMap`] internally so
+/// concurrent registers and routes never block each other.
+#[derive(Debug)]
+pub struct Demux {
+    sessions: DashMap<SessionId, tokio::sync::mpsc::Sender<InboundFrame>>,
+    waiters: DashMap<RequestId, oneshot::Sender<Result<JsonRpcResponse, ClientError>>>,
+    /// Session index for `tool.call` response waiters only. Lets the SDK
+    /// in-flight short-circuit fail every parked call for a session on a
+    /// workspace Disconnected notification without waiting for the server.
+    /// Turn-hook / session-RPC waiters are NOT indexed here, so the
+    /// short-circuit never touches them.
+    call_sessions: DashMap<RequestId, SessionId>,
+    progress: DashMap<ToolCallId, tokio::sync::mpsc::Sender<ToolCallProgressFrame>>,
+    /// Broadcast channel for connection-level notifications (no session_id).
+    notifications: tokio::sync::broadcast::Sender<Value>,
+    /// Clone of the connection's outbound sender. Used to synthesize the
+    /// overloaded (-32016) response when a session inbox is full so a
+    /// Request is rejected with an error rather than silently dropped.
+    /// `None` in unit tests that construct a bare demux.
+    outbound: Option<tokio::sync::mpsc::Sender<String>>,
+}
+
+impl Default for Demux {
+    fn default() -> Self {
+        let (notifications, _) = tokio::sync::broadcast::channel(64);
+        Self {
+            sessions: DashMap::new(),
+            waiters: DashMap::new(),
+            call_sessions: DashMap::new(),
+            progress: DashMap::new(),
+            notifications,
+            outbound: None,
+        }
+    }
+}
+
+impl Demux {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Construct a demux wired to the connection's outbound sender so the
+    /// inbox-full Request path can ship an overloaded (-32016) response.
+    pub fn with_outbound(outbound: tokio::sync::mpsc::Sender<String>) -> Self {
+        Self {
+            outbound: Some(outbound),
+            ..Self::default()
+        }
+    }
+
+    /// Subscribe to connection-level notifications (no session_id).
+    pub fn subscribe_notifications(&self) -> tokio::sync::broadcast::Receiver<Value> {
+        self.notifications.subscribe()
+    }
+
+    /// Bind `session_id` to `inbox`; replaces any existing binding.
+    /// Returns the previous sender if one existed; the caller may
+    /// drop or drain it as appropriate.
+    pub fn register_session_inbox(
+        &self,
+        session_id: SessionId,
+        inbox: tokio::sync::mpsc::Sender<InboundFrame>,
+    ) -> Option<tokio::sync::mpsc::Sender<InboundFrame>> {
+        self.sessions.insert(session_id, inbox)
+    }
+
+    /// Remove the inbox bound to `session_id`. The returned sender (if
+    /// present) is dropped by the caller, signalling EOF to its
+    /// receiver task.
+    pub fn unregister_session_inbox(
+        &self,
+        session_id: &SessionId,
+    ) -> Option<tokio::sync::mpsc::Sender<InboundFrame>> {
+        self.sessions.remove(session_id).map(|(_, sender)| sender)
+    }
+
+    /// Park a oneshot waiter for `request_id`. Crate-internal: only
+    /// the connection actor allocates request ids.
+    pub(crate) fn register_response_waiter(
+        &self,
+        request_id: RequestId,
+        waiter: oneshot::Sender<Result<JsonRpcResponse, ClientError>>,
+    ) {
+        self.waiters.insert(request_id, waiter);
+    }
+
+    /// Park a `tool.call` response waiter and record its `session_id` so the
+    /// SDK in-flight short-circuit ([`Self::fail_calls_for_session`]) can
+    /// resolve it on a workspace Disconnected notification. Crate-internal.
+    pub(crate) fn register_call_response_waiter(
+        &self,
+        request_id: RequestId,
+        session_id: SessionId,
+        waiter: oneshot::Sender<Result<JsonRpcResponse, ClientError>>,
+    ) {
+        self.call_sessions.insert(request_id.clone(), session_id);
+        self.waiters.insert(request_id, waiter);
+    }
+
+    /// Pop the waiter for `request_id`, if any. Also drops the session index
+    /// entry so the two maps stay consistent. Crate-internal.
+    pub(crate) fn take_response_waiter(
+        &self,
+        request_id: &RequestId,
+    ) -> Option<oneshot::Sender<Result<JsonRpcResponse, ClientError>>> {
+        self.call_sessions.remove(request_id);
+        self.waiters.remove(request_id).map(|(_, waiter)| waiter)
+    }
+
+    /// Fail every in-flight `tool.call` waiter bound to `session_id`,
+    /// completing each with `result_factory`. Returns the number resolved.
+    ///
+    /// Drives the SDK in-flight short-circuit: on a workspace
+    /// `ToolServerStatusChanged(Disconnected)` notification the harness fails
+    /// its parked calls for that session promptly instead of parking until
+    /// `rpc_ttl_ms`. Idempotent with the server-side cancel — each waiter is
+    /// taken at most once, so a call already resolved by the server is skipped.
+    pub(crate) fn fail_calls_for_session<F>(
+        &self,
+        session_id: &SessionId,
+        result_factory: F,
+    ) -> usize
+    where
+        F: Fn() -> ClientError,
+    {
+        // Snapshot the matching request ids first so we never hold a DashMap
+        // shard lock across the oneshot send.
+        let request_ids: Vec<RequestId> = self
+            .call_sessions
+            .iter()
+            .filter(|kv| kv.value() == session_id)
+            .map(|kv| kv.key().clone())
+            .collect();
+        let mut resolved = 0;
+        for request_id in request_ids {
+            if let Some(waiter) = self.take_response_waiter(&request_id)
+                && waiter.send(Err(result_factory())).is_ok()
+            {
+                resolved += 1;
+            }
+        }
+        resolved
+    }
+
+    /// Park a per-call progress sender keyed by `tool_call_id`.
+    /// Returns `Err(progress)` (handing the not-yet-inserted sender
+    /// back) when another in-flight call already owns the id, leaving
+    /// the prior waiter intact. The caller drops the matching
+    /// receiver to terminate the subscription — subsequent inbound
+    /// progress for the same id is silently dropped via
+    /// [`RouteOutcome::ProgressDropped`].
+    ///
+    /// Atomic check-then-insert under a single shard lock so a
+    /// concurrent caller cannot observe a transient empty slot.
+    pub(crate) fn try_register_progress_waiter(
+        &self,
+        tool_call_id: ToolCallId,
+        progress: tokio::sync::mpsc::Sender<ToolCallProgressFrame>,
+    ) -> Result<(), tokio::sync::mpsc::Sender<ToolCallProgressFrame>> {
+        use dashmap::mapref::entry::Entry;
+        match self.progress.entry(tool_call_id) {
+            Entry::Occupied(_) => Err(progress),
+            Entry::Vacant(slot) => {
+                slot.insert(progress);
+                Ok(())
+            }
+        }
+    }
+
+    /// Remove the progress sender bound to `tool_call_id`. Crate-internal;
+    /// called by the harness once the terminal frame for `tool_call_id`
+    /// has been observed.
+    pub(crate) fn unregister_progress_waiter(
+        &self,
+        tool_call_id: &ToolCallId,
+    ) -> Option<tokio::sync::mpsc::Sender<ToolCallProgressFrame>> {
+        self.progress.remove(tool_call_id).map(|(_, tx)| tx)
+    }
+
+    /// Drain every parked waiter, completing each with `result_factory`.
+    /// Used by the reconnect path to fast-fail in-flight calls with
+    /// [`ClientError::NetworkError`]. Crate-internal.
+    pub(crate) fn drain_waiters_with<F>(&self, result_factory: F)
+    where
+        F: Fn() -> ClientError,
+    {
+        // Snapshot keys, then remove individually so we never hold
+        // a DashMap shard lock across the oneshot send.
+        let keys: Vec<RequestId> = self.waiters.iter().map(|kv| kv.key().clone()).collect();
+        for key in keys {
+            if let Some((_, waiter)) = self.waiters.remove(&key) {
+                self.call_sessions.remove(&key);
+                let _ = waiter.send(Err(result_factory()));
+            }
+        }
+    }
+
+    /// Drop every parked progress sender. Used by the reconnect path
+    /// after [`Self::drain_waiters_with`]: the response waiter resolves
+    /// with `NetworkError` and the matching progress channel closes,
+    /// so any in-flight harness call stream terminates promptly
+    /// instead of stalling on a half-empty progress channel.
+    pub(crate) fn drain_progress(&self) {
+        let keys: Vec<ToolCallId> = self.progress.iter().map(|kv| kv.key().clone()).collect();
+        for key in keys {
+            self.progress.remove(&key);
+        }
+    }
+
+    /// Route a parsed JSON value. Classification rules:
+    ///
+    /// - presence of `result`/`error` → response, routed to waiter;
+    /// - method == `tool_call_progress` notification → progress waiter
+    ///   keyed by `params.tool_call_id`;
+    /// - presence of `session_id` → session inbox, request vs.
+    ///   notification distinguished by the presence of `id`;
+    /// - otherwise → [`RouteOutcome::Unrouted`].
+    ///
+    /// Routing to a session inbox or progress channel uses non-blocking
+    /// `try_send`. A full inbox or progress channel returns the matching
+    /// `*Full` variant; a dropped receiver returns the matching
+    /// `*Dropped` variant. Either way the frame is dropped without
+    /// awaiting the consumer, so a slow handler never starves other
+    /// sessions or calls multiplexed onto the same connection.
+    pub fn route(&self, frame: Value) -> RouteOutcome {
+        crate::metrics::demux_inbox_depth_set(self.sessions.len() as i64);
+        if frame.get("result").is_some() || frame.get("error").is_some() {
+            return self.route_response(frame);
+        }
+        if frame.get("method").and_then(Value::as_str) == Some("tool_call_progress") {
+            return self.route_progress(frame);
+        }
+        if frame.get("session_id").is_some() {
+            return self.route_session(frame);
+        }
+        // Connection-level notification (e.g. session.bind, session.unbind).
+        if frame.get("method").is_some() {
+            let _ = self.notifications.send(frame);
+            return RouteOutcome::Notification;
+        }
+        RouteOutcome::Unrouted
+    }
+
+    fn route_progress(&self, frame: Value) -> RouteOutcome {
+        let Some(params) = frame.get("params") else {
+            return RouteOutcome::Unrouted;
+        };
+        let Some(call_id_str) = params.get("tool_call_id").and_then(Value::as_str) else {
+            return RouteOutcome::Unrouted;
+        };
+        let Ok(tool_call_id) = ToolCallId::new(call_id_str) else {
+            return RouteOutcome::Unrouted;
+        };
+        let Some(sender) = self.progress.get(&tool_call_id) else {
+            return RouteOutcome::UnknownProgress;
+        };
+        let tx = sender.value().clone();
+        drop(sender);
+        let progress_frame: ToolCallProgressFrame = match serde_json::from_value(params.clone()) {
+            Ok(p) => p,
+            Err(err) => {
+                warn!(%tool_call_id, ?err, "failed to decode tool_call_progress params");
+                return RouteOutcome::Unrouted;
+            }
+        };
+        match tx.try_send(progress_frame) {
+            Ok(()) => RouteOutcome::Progress,
+            Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
+                warn!(%tool_call_id, "progress channel full; dropping inbound progress frame");
+                RouteOutcome::ProgressFull
+            }
+            Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
+                self.progress.remove(&tool_call_id);
+                RouteOutcome::ProgressDropped
+            }
+        }
+    }
+
+    fn route_response(&self, frame: Value) -> RouteOutcome {
+        let Some(id_value) = frame.get("id") else {
+            return RouteOutcome::Unrouted;
+        };
+        let request_id = match id_value {
+            Value::String(s) => RequestId::new(s.as_str()).ok(),
+            Value::Number(n) => RequestId::new(n.to_string()).ok(),
+            _ => None,
+        };
+        let Some(request_id) = request_id else {
+            return RouteOutcome::Unrouted;
+        };
+        let Some(waiter) = self.take_response_waiter(&request_id) else {
+            return RouteOutcome::Unrouted;
+        };
+        let parsed: Result<JsonRpcResponse, ClientError> =
+            serde_json::from_value::<JsonRpcResponse>(frame).map_err(ClientError::from);
+        let _ = waiter.send(parsed);
+        RouteOutcome::Response
+    }
+
+    fn route_session(&self, frame: Value) -> RouteOutcome {
+        let Some(sid_str) = frame.get("session_id").and_then(Value::as_str) else {
+            return RouteOutcome::Unrouted;
+        };
+        let Ok(session_id) = SessionId::new(sid_str) else {
+            return RouteOutcome::Unrouted;
+        };
+        let Some(sender) = self.sessions.get(&session_id) else {
+            return RouteOutcome::UnknownSession;
+        };
+        let inbox = sender.value().clone();
+        drop(sender);
+        let kind = if frame.get("id").is_some() {
+            InboundFrame::Request(frame)
+        } else {
+            InboundFrame::Notification(frame)
+        };
+        match inbox.try_send(kind) {
+            Ok(()) => RouteOutcome::Session,
+            Err(tokio::sync::mpsc::error::TrySendError::Full(frame)) => {
+                self.reject_inbox_full(&session_id, frame);
+                RouteOutcome::InboxFull
+            }
+            Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
+                warn!(%session_id, "session inbox dropped; binding stale");
+                self.sessions.remove(&session_id);
+                RouteOutcome::SessionDropped
+            }
+        }
+    }
+
+    /// Handle a full session inbox without blocking the reader.
+    ///
+    /// A Request (has an `id`) is rejected with the shared overloaded
+    /// (-32016 "tool_busy") response on a best-effort `try_send`; if the
+    /// outbound is *also* full the rejection itself is dropped and metered
+    /// (`inbox_full_reject_send_failed`). A Notification (no `id`) stays
+    /// fire-and-forget and is metered (`inbox_full_notification_dropped`).
+    fn reject_inbox_full(&self, session_id: &SessionId, frame: InboundFrame) {
+        let InboundFrame::Request(value) = frame else {
+            crate::metrics::inbox_full_notification_dropped();
+            return;
+        };
+        crate::metrics::inbox_full_request_rejected();
+        warn!(%session_id, "session inbox full; rejecting request with tool_busy");
+        let Some(out) = &self.outbound else {
+            return;
+        };
+        // A `Request` always carries an `id` (that is how `route_session`
+        // classifies it). A well-formed id deserializes into a `JsonRpcId`;
+        // a malformed id (object/array/bool/null) cannot, but the request
+        // must STILL get an overloaded response rather than be silently
+        // dropped, so we fall back to echoing the raw id JSON as a string.
+        let raw_id = value.get("id");
+        let id = raw_id
+            .and_then(|v| serde_json::from_value::<JsonRpcId>(v.clone()).ok())
+            .unwrap_or_else(|| {
+                JsonRpcId::new_string(raw_id.map(ToString::to_string).unwrap_or_default())
+            });
+        let response = crate::admission::overloaded_response(id, session_id.clone());
+        let Ok(text) = serde_json::to_string(&response) else {
+            return;
+        };
+        if out.try_send(text).is_err() {
+            crate::metrics::inbox_full_reject_send_failed();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+    use tokio::sync::mpsc;
+
+    #[tokio::test]
+    async fn response_route_matches_waiter() {
+        let demux = Demux::new();
+        let request_id = RequestId::new("r1").expect("valid");
+        let (tx, rx) = oneshot::channel();
+        demux.register_response_waiter(request_id.clone(), tx);
+        let outcome = demux.route(json!({
+            "jsonrpc": "2.0",
+            "id": "r1",
+            "result": {"outcome": "bound"},
+        }));
+        assert_eq!(outcome, RouteOutcome::Response);
+        let resp = rx.await.expect("waiter").expect("ok");
+        assert_eq!(resp.id.to_string(), "r1");
+    }
+
+    #[tokio::test]
+    async fn fail_calls_for_session_resolves_only_matching_call_waiters() {
+        // Fails exactly the session's `tool.call` waiters; other sessions'
+        // calls and non-call (turn-hook) waiters stay parked.
+        let demux = Demux::new();
+        let s1 = SessionId::new("s1").expect("valid");
+        let s2 = SessionId::new("s2").expect("valid");
+
+        let (tx_a, rx_a) = oneshot::channel();
+        let (tx_b, rx_b) = oneshot::channel();
+        let (tx_other, rx_other) = oneshot::channel();
+        // Two calls on s1, one on s2.
+        demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a);
+        demux.register_call_response_waiter(RequestId::new("b").unwrap(), s1.clone(), tx_b);
+        demux.register_call_response_waiter(RequestId::new("c").unwrap(), s2.clone(), tx_other);
+        // A non-call waiter (e.g. a turn hook) on s1 — NOT session-indexed.
+        let (tx_hook, rx_hook) = oneshot::channel();
+        demux.register_response_waiter(RequestId::new("hook").unwrap(), tx_hook);
+
+        let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned()));
+        assert_eq!(n, 2, "only the two s1 call waiters are failed");
+
+        assert!(matches!(rx_a.await, Ok(Err(ClientError::NetworkError(_)))));
+        assert!(matches!(rx_b.await, Ok(Err(ClientError::NetworkError(_)))));
+        // s2's call and the turn-hook waiter are untouched (still parked).
+        assert!(
+            demux
+                .take_response_waiter(&RequestId::new("c").unwrap())
+                .is_some(),
+            "the s2 call must remain parked"
+        );
+        assert!(
+            demux
+                .take_response_waiter(&RequestId::new("hook").unwrap())
+                .is_some(),
+            "the turn-hook waiter must remain parked"
+        );
+        // Keep the receivers alive until the asserts above ran.
+        drop((rx_other, rx_hook));
+    }
+
+    #[tokio::test]
+    async fn fail_calls_for_session_is_idempotent_after_resolution() {
+        // A call already resolved (waiter taken) must not be double-counted by
+        // the short-circuit.
+        let demux = Demux::new();
+        let s1 = SessionId::new("s1").expect("valid");
+        let (tx_a, rx_a) = oneshot::channel();
+        demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a);
+        // Simulate the server-side resolution taking the waiter first.
+        let waiter = demux
+            .take_response_waiter(&RequestId::new("a").unwrap())
+            .expect("waiter present");
+        drop(waiter);
+        drop(rx_a);
+        let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned()));
+        assert_eq!(n, 0, "already-resolved call is not re-failed");
+    }
+
+    #[tokio::test]
+    async fn short_circuit_then_late_response_is_unrouted() {
+        // A short-circuit that resolves first leaves no waiter, so a late server
+        // response for the same id is dropped (no double-resolve).
+        let demux = Demux::new();
+        let s1 = SessionId::new("s1").expect("valid");
+        let (tx_a, rx_a) = oneshot::channel();
+        demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a);
+
+        let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned()));
+        assert_eq!(n, 1);
+        assert!(matches!(rx_a.await, Ok(Err(ClientError::NetworkError(_)))));
+
+        let outcome = demux.route(json!({ "jsonrpc": "2.0", "id": "a", "result": {} }));
+        assert_eq!(
+            outcome,
+            RouteOutcome::Unrouted,
+            "the late normal response must not double-resolve the call"
+        );
+    }
+
+    #[tokio::test]
+    async fn session_route_pushes_to_inbox() {
+        let demux = Demux::new();
+        let session = SessionId::new("s1").expect("valid");
+        let (tx, mut rx) = mpsc::channel(4);
+        demux.register_session_inbox(session.clone(), tx);
+        let frame = json!({
+            "jsonrpc": "2.0",
+            "id": "x",
+            "session_id": "s1",
+            "method": "tool_call_request",
+            "params": {},
+        });
+        let outcome = demux.route(frame.clone());
+        assert_eq!(outcome, RouteOutcome::Session);
+        match rx.recv().await {
+            Some(InboundFrame::Request(value)) => assert_eq!(value, frame),
+            other => panic!("expected request inbound; got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn reverse_hook_request_routes_to_inbox_as_request() {
+        // A reverse hook request carries an `id`, so it must route to the
+        // inbox as `Request` (not `Notification`) for the harness to answer.
+        let demux = Demux::new();
+        let session = SessionId::new("s1").expect("valid");
+        let (tx, mut rx) = mpsc::channel(4);
+        demux.register_session_inbox(session.clone(), tx);
+        let hook = xai_tool_protocol::HookFrame::custom_request(
+            session.clone(),
+            "hook-7".to_owned(),
+            crate::harness::PERMISSION_REQUEST_KIND.to_owned(),
+            json!({}),
+        );
+        let frame = json!({
+            "jsonrpc": "2.0",
+            "id": "h1",
+            "session_id": "s1",
+            "method": xai_tool_protocol::Method::Hook.as_wire_str(),
+            "params": serde_json::to_value(&hook).expect("serialize hook"),
+        });
+        assert_eq!(demux.route(frame.clone()), RouteOutcome::Session);
+        match rx.recv().await {
+            Some(InboundFrame::Request(value)) => assert_eq!(value, frame),
+            other => panic!("expected request inbound; got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn notification_classified_without_id() {
+        let demux = Demux::new();
+        let session = SessionId::new("s1").expect("valid");
+        let (tx, mut rx) = mpsc::channel(4);
+        demux.register_session_inbox(session.clone(), tx);
+        let frame = json!({
+            "jsonrpc": "2.0",
+            "session_id": "s1",
+            "method": "tool.notification",
+            "params": {},
+        });
+        let outcome = demux.route(frame);
+        assert_eq!(outcome, RouteOutcome::Session);
+        match rx.recv().await {
+            Some(InboundFrame::Notification(_)) => {}
+            other => panic!("expected notification; got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn unknown_session_returns_unknown_session() {
+        let demux = Demux::new();
+        let outcome = demux.route(
+            json!({"jsonrpc":"2.0","id":"x","session_id":"missing","method":"x","params":{}}),
+        );
+        assert_eq!(outcome, RouteOutcome::UnknownSession);
+    }
+
+    #[tokio::test]
+    async fn unknown_request_id_returns_unrouted() {
+        let demux = Demux::new();
+        let outcome = demux.route(json!({
+            "jsonrpc": "2.0",
+            "id": "missing",
+            "result": {},
+        }));
+        assert_eq!(outcome, RouteOutcome::Unrouted);
+    }
+
+    #[tokio::test]
+    async fn full_inbox_returns_inbox_full_without_blocking() {
+        let demux = Demux::new();
+        let session = SessionId::new("backed_up").expect("valid");
+        let (tx, _rx) = mpsc::channel(1);
+        demux.register_session_inbox(session.clone(), tx);
+        let frame = || {
+            json!({
+                "jsonrpc": "2.0",
+                "id": "x",
+                "session_id": "backed_up",
+                "method": "tool_call_request",
+                "params": {},
+            })
+        };
+        // First send fills capacity.
+        assert_eq!(demux.route(frame()), RouteOutcome::Session);
+        // Second send must NOT block; it returns InboxFull.
+        assert_eq!(demux.route(frame()), RouteOutcome::InboxFull);
+    }
+
+    #[tokio::test]
+    async fn dropped_receiver_returns_session_dropped() {
+        let demux = Demux::new();
+        let session = SessionId::new("gone").expect("valid");
+        let (tx, rx) = mpsc::channel(1);
+        demux.register_session_inbox(session.clone(), tx);
+        drop(rx);
+        let frame = json!({
+            "jsonrpc": "2.0",
+            "id": "x",
+            "session_id": "gone",
+            "method": "tool_call_request",
+            "params": {},
+        });
+        assert_eq!(demux.route(frame), RouteOutcome::SessionDropped);
+        // Stale binding should have been removed.
+        assert!(demux.sessions.get(&session).is_none());
+    }
+
+    #[tokio::test]
+    async fn inbox_full_request_synthesizes_overloaded_response_onto_outbound() {
+        // A full session inbox for a Request must produce the shared
+        // -32016 "tool_busy" response on outbound, not a silent drop.
+        let (out_tx, mut out_rx) = mpsc::channel::<String>(4);
+        let demux = Demux::with_outbound(out_tx);
+        let session = SessionId::new("busy").expect("valid");
+        let (tx, _rx) = mpsc::channel(1);
+        demux.register_session_inbox(session.clone(), tx);
+        let frame = |id: &str| {
+            json!({
+                "jsonrpc": "2.0",
+                "id": id,
+                "session_id": "busy",
+                "method": "tool_call_request",
+                "params": {},
+            })
+        };
+        // First fills capacity (cap 1); second overflows → InboxFull.
+        assert_eq!(demux.route(frame("a")), RouteOutcome::Session);
+        assert_eq!(demux.route(frame("b")), RouteOutcome::InboxFull);
+
+        let text = out_rx.try_recv().expect("overloaded response enqueued");
+        let wire: Value = serde_json::from_str(&text).expect("valid json");
+        assert_eq!(wire["id"], "b");
+        assert_eq!(wire["session_id"], "busy");
+        assert_eq!(wire["error"]["code"], -32016);
+        assert_eq!(wire["error"]["data"]["code"], "tool_busy");
+        assert_eq!(wire["error"]["data"]["retryable"], true);
+        assert!(
+            out_rx.try_recv().is_err(),
+            "exactly one rejection emitted for one overflow"
+        );
+    }
+
+    #[tokio::test]
+    async fn inbox_full_request_with_malformed_id_still_emits_overloaded_response() {
+        // A Request whose `id` is present but not a valid JsonRpcId
+        // (object/array/null) must NOT be silently dropped on a full
+        // inbox: it still gets the shared -32016 response, with the raw
+        // id echoed back as a string.
+        let (out_tx, mut out_rx) = mpsc::channel::<String>(4);
+        let demux = Demux::with_outbound(out_tx);
+        let session = SessionId::new("bad_id").expect("valid");
+        let (tx, _rx) = mpsc::channel(1);
+        demux.register_session_inbox(session.clone(), tx);
+        let frame = |id: Value| {
+            json!({
+                "jsonrpc": "2.0",
+                "id": id,
+                "session_id": "bad_id",
+                "method": "tool_call_request",
+                "params": {},
+            })
+        };
+        // First fills capacity (cap 1); the malformed-id second overflows.
+        assert_eq!(demux.route(frame(json!("a"))), RouteOutcome::Session);
+        assert_eq!(
+            demux.route(frame(json!({ "nested": 1 }))),
+            RouteOutcome::InboxFull
+        );
+
+        let text = out_rx.try_recv().expect("overloaded response enqueued");
+        let wire: Value = serde_json::from_str(&text).expect("valid json");
+        assert_eq!(
+            wire["id"], "{\"nested\":1}",
+            "malformed id is echoed back as its raw JSON text"
+        );
+        assert_eq!(wire["error"]["code"], -32016);
+        assert_eq!(wire["error"]["data"]["code"], "tool_busy");
+    }
+
+    #[tokio::test]
+    async fn inbox_full_notification_is_dropped_without_outbound_response() {
+        // A Notification (no id) on a full inbox stays fire-and-forget:
+        // no synthesized response is emitted.
+        let (out_tx, mut out_rx) = mpsc::channel::<String>(4);
+        let demux = Demux::with_outbound(out_tx);
+        let session = SessionId::new("notif_busy").expect("valid");
+        let (tx, _rx) = mpsc::channel(1);
+        demux.register_session_inbox(session.clone(), tx);
+        let notif = || {
+            json!({
+                "jsonrpc": "2.0",
+                "session_id": "notif_busy",
+                "method": "tool.notification",
+                "params": {},
+            })
+        };
+        // First notification fills capacity; second overflows.
+        assert_eq!(demux.route(notif()), RouteOutcome::Session);
+        assert_eq!(demux.route(notif()), RouteOutcome::InboxFull);
+        assert!(
+            out_rx.try_recv().is_err(),
+            "notifications must not synthesize an outbound response"
+        );
+    }
+
+    #[tokio::test]
+    async fn inbox_full_request_without_outbound_does_not_panic() {
+        // A bare demux (no outbound, e.g. unit context) must still report
+        // InboxFull cleanly when it cannot synthesize a rejection.
+        let demux = Demux::new();
+        let session = SessionId::new("no_out").expect("valid");
+        let (tx, _rx) = mpsc::channel(1);
+        demux.register_session_inbox(session.clone(), tx);
+        let frame = || {
+            json!({
+                "jsonrpc": "2.0",
+                "id": "x",
+                "session_id": "no_out",
+                "method": "tool_call_request",
+                "params": {},
+            })
+        };
+        assert_eq!(demux.route(frame()), RouteOutcome::Session);
+        assert_eq!(demux.route(frame()), RouteOutcome::InboxFull);
+    }
+
+    #[tokio::test]
+    async fn progress_route_pushes_to_progress_waiter() {
+        let demux = Demux::new();
+        let call_id = ToolCallId::new_v7();
+        let (tx, mut rx) = mpsc::channel(4);
+        demux
+            .try_register_progress_waiter(call_id.clone(), tx)
+            .expect("first registration");
+        let frame = json!({
+            "jsonrpc": "2.0",
+            "session_id": "any",
+            "method": "tool_call_progress",
+            "params": {
+                "tool_call_id": call_id.as_str(),
+                "kind": "log_chunk",
+                "body": {"text": "hello"},
+            },
+        });
+        let outcome = demux.route(frame);
+        assert_eq!(outcome, RouteOutcome::Progress);
+        let progress = rx.recv().await.expect("progress frame");
+        assert_eq!(progress.tool_call_id, call_id);
+        assert_eq!(progress.kind, "log_chunk");
+        assert_eq!(progress.body, json!({"text": "hello"}));
+    }
+
+    #[tokio::test]
+    async fn progress_with_no_waiter_returns_unknown_progress() {
+        let demux = Demux::new();
+        let call_id = ToolCallId::new_v7();
+        let frame = json!({
+            "jsonrpc": "2.0",
+            "session_id": "any",
+            "method": "tool_call_progress",
+            "params": {
+                "tool_call_id": call_id.as_str(),
+                "kind": "log_chunk",
+                "body": {},
+            },
+        });
+        assert_eq!(demux.route(frame), RouteOutcome::UnknownProgress);
+    }
+
+    #[tokio::test]
+    async fn dropped_progress_receiver_returns_progress_dropped() {
+        let demux = Demux::new();
+        let call_id = ToolCallId::new_v7();
+        let (tx, rx) = mpsc::channel(1);
+        demux
+            .try_register_progress_waiter(call_id.clone(), tx)
+            .expect("first registration");
+        drop(rx);
+        let frame = json!({
+            "jsonrpc": "2.0",
+            "session_id": "any",
+            "method": "tool_call_progress",
+            "params": {
+                "tool_call_id": call_id.as_str(),
+                "kind": "x",
+                "body": {},
+            },
+        });
+        assert_eq!(demux.route(frame), RouteOutcome::ProgressDropped);
+        assert!(demux.progress.get(&call_id).is_none());
+    }
+
+    #[tokio::test]
+    async fn unregister_progress_waiter_returns_sender_when_present() {
+        let demux = Demux::new();
+        let call_id = ToolCallId::new_v7();
+        let (tx, _rx) = mpsc::channel::<ToolCallProgressFrame>(1);
+        demux
+            .try_register_progress_waiter(call_id.clone(), tx)
+            .expect("first registration");
+        assert!(demux.unregister_progress_waiter(&call_id).is_some());
+        assert!(demux.unregister_progress_waiter(&call_id).is_none());
+    }
+
+    #[tokio::test]
+    async fn try_register_progress_waiter_rejects_collision_and_preserves_existing() {
+        let demux = Demux::new();
+        let call_id = ToolCallId::new_v7();
+        let (tx_first, mut rx_first) = mpsc::channel::<ToolCallProgressFrame>(1);
+        demux
+            .try_register_progress_waiter(call_id.clone(), tx_first)
+            .expect("first registration");
+        let (tx_second, _rx_second) = mpsc::channel::<ToolCallProgressFrame>(1);
+        let returned = demux
+            .try_register_progress_waiter(call_id.clone(), tx_second)
+            .expect_err("collision returns the rejected sender");
+        // Returned sender is independent of the live one: dropping
+        // it must not close the original receiver.
+        drop(returned);
+        let frame = json!({
+            "jsonrpc": "2.0",
+            "session_id": "any",
+            "method": "tool_call_progress",
+            "params": {
+                "tool_call_id": call_id.as_str(),
+                "kind": "log_chunk",
+                "body": {"text": "first"},
+            },
+        });
+        assert_eq!(demux.route(frame), RouteOutcome::Progress);
+        let progress = rx_first.recv().await.expect("original receiver still live");
+        assert_eq!(progress.body, json!({"text": "first"}));
+    }
+
+    #[tokio::test]
+    async fn full_progress_channel_returns_progress_full_without_blocking() {
+        let demux = Demux::new();
+        let call_id = ToolCallId::new_v7();
+        let (tx, _rx) = mpsc::channel::<ToolCallProgressFrame>(1);
+        demux
+            .try_register_progress_waiter(call_id.clone(), tx)
+            .expect("first registration");
+        let frame = || {
+            json!({
+                "jsonrpc": "2.0",
+                "session_id": "any",
+                "method": "tool_call_progress",
+                "params": {
+                    "tool_call_id": call_id.as_str(),
+                    "kind": "x",
+                    "body": {},
+                },
+            })
+        };
+        // First send fills capacity (mpsc(1)).
+        assert_eq!(demux.route(frame()), RouteOutcome::Progress);
+        // Second send must NOT block; it returns ProgressFull.
+        assert_eq!(demux.route(frame()), RouteOutcome::ProgressFull);
+    }
+
+    #[tokio::test]
+    async fn drain_progress_removes_all_waiters_and_drops_senders() {
+        let demux = Demux::new();
+        let call_a = ToolCallId::new_v7();
+        let call_b = ToolCallId::new_v7();
+        let (tx_a, mut rx_a) = mpsc::channel::<ToolCallProgressFrame>(1);
+        let (tx_b, mut rx_b) = mpsc::channel::<ToolCallProgressFrame>(1);
+        demux
+            .try_register_progress_waiter(call_a.clone(), tx_a)
+            .expect("first registration");
+        demux
+            .try_register_progress_waiter(call_b.clone(), tx_b)
+            .expect("first registration");
+        assert_eq!(demux.progress.len(), 2);
+
+        demux.drain_progress();
+
+        // Post-drain: every entry removed.
+        assert_eq!(demux.progress.len(), 0);
+        assert!(demux.progress.get(&call_a).is_none());
+        assert!(demux.progress.get(&call_b).is_none());
+        // The senders held by the demux were dropped, so each
+        // receiver sees `None` (channel closed).
+        assert!(
+            rx_a.recv().await.is_none(),
+            "sender dropped → receiver closes"
+        );
+        assert!(
+            rx_b.recv().await.is_none(),
+            "sender dropped → receiver closes"
+        );
+    }
+}

crates/common/xai-computer-hub-sdk/src/donate_pump.rs

@@ -0,0 +1,206 @@
+//! Shared donation transport: a bounded retry buffer + in-order drain
+//! barrier, parameterized over a `donate` closure. Traces, logs, and
+//! metrics all pump through this; failed sends are retained briefly,
+//! overflow drops payloads — telemetry, never correctness.
+
+use std::collections::VecDeque;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue, any_value};
+use opentelemetry_proto::tonic::resource::v1::Resource;
+use tokio::sync::{mpsc, oneshot};
+
+/// Bound on payloads queued before the pump drains them.
+pub(crate) const PENDING_FLUSHES: usize = 8;
+/// Payloads retained across failed sends (disconnect/reconnect window).
+pub(crate) const RETRY_CAP: usize = 8;
+
+// ---------------------------------------------------------------------------
+// Shared OTLP encoding helpers
+//
+// Reused by the log and metric donation clients so the AnyValue/KeyValue/
+// Resource construction lives in one place instead of being copy-pasted per
+// client. (`trace_donate` builds its payload via `opentelemetry_sdk`'s own
+// conversion and does not use these.)
+// ---------------------------------------------------------------------------
+
+/// Current wall-clock time as Unix-epoch nanoseconds (OTLP `time_unix_nano`).
+pub(crate) fn now_unix_nanos() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_nanos() as u64)
+        .unwrap_or(0)
+}
+
+/// OTLP string `AnyValue`.
+pub(crate) fn string_value(s: String) -> AnyValue {
+    AnyValue {
+        value: Some(any_value::Value::StringValue(s)),
+    }
+}
+
+/// OTLP string-valued `KeyValue`.
+pub(crate) fn string_kv(key: &str, value: String) -> KeyValue {
+    KeyValue {
+        key: key.to_owned(),
+        value: Some(string_value(value)),
+        ..Default::default()
+    }
+}
+
+/// OTLP `Resource` carrying just `service.name`.
+pub(crate) fn make_resource(service_name: String) -> Resource {
+    Resource {
+        attributes: vec![string_kv("service.name", service_name)],
+        ..Default::default()
+    }
+}
+
+pub(crate) enum PumpMsg {
+    /// Base64 OTLP request, ready for the wire.
+    Payload(String),
+    /// In-order drain fence — a barrier, not a timeout.
+    Barrier(oneshot::Sender<()>),
+}
+
+/// Resolves once every payload queued before this call has had a send
+/// attempt. Call after the producer's flush (e.g. `fastrace::flush()`).
+pub(crate) async fn drain_via(tx: &mpsc::Sender<PumpMsg>) {
+    let (ack_tx, ack_rx) = oneshot::channel();
+    if tx.send(PumpMsg::Barrier(ack_tx)).await.is_ok() {
+        let _ = ack_rx.await;
+    }
+}
+
+/// `donate` hands the payload back so a failed send retains it
+/// without cloning.
+pub(crate) async fn run_pump<D, F>(mut rx: mpsc::Receiver<PumpMsg>, donate: D)
+where
+    D: Fn(String) -> F,
+    F: std::future::Future<Output = (bool, String)>,
+{
+    let mut retry: VecDeque<String> = VecDeque::new();
+    while let Some(msg) = rx.recv().await {
+        match msg {
+            PumpMsg::Payload(payload) => {
+                if retry.len() == RETRY_CAP {
+                    retry.pop_front();
+                    tracing::debug!("donation retry buffer full; dropping oldest payload");
+                }
+                retry.push_back(payload);
+            }
+            PumpMsg::Barrier(ack) => {
+                attempt_sends(&mut retry, &donate).await;
+                let _ = ack.send(());
+                continue;
+            }
+        }
+        attempt_sends(&mut retry, &donate).await;
+    }
+}
+
+/// Send in order, stopping at the first failure; the remainder stays
+/// queued for the next wake.
+async fn attempt_sends<D, F>(retry: &mut VecDeque<String>, donate: &D)
+where
+    D: Fn(String) -> F,
+    F: std::future::Future<Output = (bool, String)>,
+{
+    while let Some(payload) = retry.pop_front() {
+        let (ok, payload) = donate(payload).await;
+        if !ok {
+            tracing::debug!("donation send failed; retaining payload for retry");
+            retry.push_front(payload);
+            break;
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicBool, Ordering};
+
+    use parking_lot::Mutex;
+
+    use super::*;
+
+    fn payload(tag: u64) -> PumpMsg {
+        PumpMsg::Payload(format!("payload-{tag}"))
+    }
+
+    /// The drain barrier acks even while the link is down.
+    #[tokio::test]
+    async fn pump_retries_failed_payloads_across_reconnect() {
+        let healthy = Arc::new(AtomicBool::new(false));
+        let sent: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+        let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
+        let pump = {
+            let healthy = Arc::clone(&healthy);
+            let sent = Arc::clone(&sent);
+            tokio::spawn(run_pump(rx, move |p: String| {
+                let healthy = Arc::clone(&healthy);
+                let sent = Arc::clone(&sent);
+                async move {
+                    if healthy.load(Ordering::SeqCst) {
+                        sent.lock().push(p.clone());
+                        (true, p)
+                    } else {
+                        (false, p)
+                    }
+                }
+            }))
+        };
+
+        tx.send(payload(1)).await.unwrap();
+        tx.send(payload(2)).await.unwrap();
+        drain_via(&tx).await;
+        assert!(sent.lock().is_empty(), "nothing sent while link is down");
+
+        healthy.store(true, Ordering::SeqCst);
+        drain_via(&tx).await;
+        assert_eq!(*sent.lock(), vec!["payload-1", "payload-2"]);
+
+        drop(tx);
+        pump.await.expect("pump must exit cleanly");
+    }
+
+    #[tokio::test]
+    async fn pump_retry_buffer_drops_oldest_beyond_cap() {
+        let sent: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+        let healthy = Arc::new(AtomicBool::new(false));
+        let (tx, rx) = mpsc::channel::<PumpMsg>(RETRY_CAP + 2);
+        let pump = {
+            let healthy = Arc::clone(&healthy);
+            let sent = Arc::clone(&sent);
+            tokio::spawn(run_pump(rx, move |p: String| {
+                let healthy = Arc::clone(&healthy);
+                let sent = Arc::clone(&sent);
+                async move {
+                    if healthy.load(Ordering::SeqCst) {
+                        sent.lock().push(p.clone());
+                        (true, p)
+                    } else {
+                        (false, p)
+                    }
+                }
+            }))
+        };
+
+        for i in 0..=(RETRY_CAP as u64) {
+            tx.send(payload(i + 1)).await.unwrap();
+        }
+        drain_via(&tx).await;
+
+        healthy.store(true, Ordering::SeqCst);
+        drain_via(&tx).await;
+        {
+            let sent = sent.lock();
+            assert_eq!(sent.len(), RETRY_CAP, "buffer bounded at RETRY_CAP");
+            assert_eq!(sent[0], "payload-2", "oldest payload evicted first");
+        }
+
+        drop(tx);
+        pump.await.expect("pump must exit cleanly");
+    }
+}

crates/common/xai-computer-hub-sdk/src/error.rs

@@ -0,0 +1,346 @@
+//! Client-side error taxonomy.
+//!
+//! Wire-level [`xai_tool_protocol::ToolErrorWire`] variants and JSON-RPC
+//! error envelopes are mapped into the smaller [`ClientError`] vocabulary
+//! at the SDK boundary so consumers can match on a single enum without
+//! re-deriving the numeric/string code mapping.
+
+use thiserror::Error;
+use url::Url;
+use xai_tool_protocol::{IdError, JsonRpcError, ToolCallId, ToolErrorWire};
+
+/// Errors surfaced by the client SDK.
+#[derive(Debug, Error)]
+pub enum ClientError {
+    /// WebSocket transport failure: failed to connect, dropped socket,
+    /// or in-flight request interrupted by a reconnect cycle.
+    #[error("network error: {0}")]
+    NetworkError(String),
+
+    /// Wire-protocol violation: malformed JSON, unexpected method,
+    /// hello/hello_ack mismatch, or unsupported `protocol_version`.
+    #[error("protocol error: {0}")]
+    ProtocolError(String),
+
+    /// Authentication or authorisation rejected by the server.
+    #[error("auth error: {0}")]
+    AuthError(String),
+
+    /// Server rejected the WebSocket upgrade with an HTTP auth status
+    /// (401/403). Non-retryable: replaying the same credential is
+    /// rejected identically, so the reconnect loop classifies this as
+    /// fatal instead of retrying forever.
+    #[error("handshake auth failed: HTTP {status}")]
+    HandshakeAuthFailed { status: u16 },
+
+    /// `register_tool` / `register_session` ack reported a conflict
+    /// (cross-connection contention or an already-bound entry the
+    /// caller did not expect).
+    #[error("registration conflict: {0}")]
+    RegistrationConflict(String),
+
+    /// Outbound mpsc full or call-site bounded wait elapsed before the
+    /// frame could be enqueued. Distinct from [`Self::NetworkError`]:
+    /// the socket may still be healthy.
+    #[error("backpressure: {0}")]
+    BackpressureError(String),
+
+    /// JSON serialise / deserialise failure inside the SDK.
+    #[error("serde error: {0}")]
+    Serde(String),
+
+    /// Builder consistency error: missing URL, missing auth, etc.
+    #[error("invalid configuration: {0}")]
+    InvalidConfig(String),
+
+    /// Wrapped wire-format tool error; surfaces the upstream
+    /// [`ToolErrorWire`] variant verbatim for callers that need to
+    /// switch on the stable string code.
+    #[error(transparent)]
+    Wire(ToolErrorWire),
+
+    /// Server-side close / shutdown signal received during steady state.
+    #[error("server closed connection: {0}")]
+    Closed(String),
+
+    /// Refused to send credentials over an insecure `ws://` scheme to a
+    /// non-loopback host. Local-loopback (`127.0.0.1`, `::1`,
+    /// `localhost`) is the only exception; every other host MUST be
+    /// reached over `wss://` so the bearer token never crosses the
+    /// network in plaintext.
+    #[error(
+        "insecure scheme: refusing to send credentials over plaintext ws:// to non-loopback host {url}"
+    )]
+    InsecureScheme { url: Url },
+
+    /// Caller passed a `ToolCallId` that already keys an in-flight
+    /// dispatch on the same connection. The prior call's progress
+    /// waiter and response correlation are left intact; this error
+    /// surfaces synchronously so the second caller can retry with a
+    /// fresh id. Mint a fresh [`ToolCallId::new_v7`] (or use
+    /// [`xai_tool_runtime::ToolCallContext::default`], which does so)
+    /// per call. This is client misuse, not a transport or server
+    /// failure.
+    #[error("call_id {call_id} already in flight on this connection")]
+    CallIdInUse { call_id: ToolCallId },
+}
+
+impl ClientError {
+    /// Map a JSON-RPC envelope error into a [`ClientError`]. The
+    /// envelope's `data` payload (when present) carries the stable
+    /// [`ToolErrorWire`] discriminator; the numeric `code` is used as a
+    /// coarse fallback when `data` is absent or undecodable.
+    pub fn from_jsonrpc_error(err: JsonRpcError) -> Self {
+        if let Some(data) = err.data
+            && let Ok(wire) = serde_json::from_value::<ToolErrorWire>(data)
+        {
+            return Self::from_wire(wire);
+        }
+        match err.code {
+            -32002 | -32003 => Self::AuthError(err.message),
+            -32004 => Self::NetworkError(err.message),
+            -32600..=-32500 => Self::ProtocolError(err.message),
+            _ => Self::Wire(ToolErrorWire::Custom {
+                subcode: format!("jsonrpc_{}", err.code),
+                message: err.message,
+                details: None,
+            }),
+        }
+    }
+
+    /// `true` when a `data`-less envelope collapsed to the given `jsonrpc_<code>`
+    /// subcode (see [`Self::from_jsonrpc_error`]); shared by the bind recognizers.
+    fn has_collapsed_jsonrpc_subcode(&self, subcode: &str) -> bool {
+        matches!(
+            self,
+            Self::Wire(ToolErrorWire::Custom { subcode: s, .. }) if s == subcode
+        )
+    }
+
+    /// `true` for the server's "server not found" bind rejection (JSON-RPC `-32601`):
+    /// no workspace-server is registered for this user.
+    pub fn is_server_not_found(&self) -> bool {
+        self.has_collapsed_jsonrpc_subcode("jsonrpc_-32601")
+    }
+
+    /// `true` for the server's `-32013` "server found but bind did not complete" error
+    /// (the `ServerBindOutcome::Unavailable` cases). Recognized so the harness
+    /// re-provisions this recoverable case, distinct from [`Self::is_server_not_found`].
+    pub fn is_tool_unavailable(&self) -> bool {
+        self.has_collapsed_jsonrpc_subcode("jsonrpc_-32013")
+    }
+
+    /// Map a [`ToolErrorWire`] variant into the SDK error taxonomy.
+    pub fn from_wire(wire: ToolErrorWire) -> Self {
+        match wire {
+            ToolErrorWire::PermissionDenied { reason } => Self::AuthError(reason),
+            ToolErrorWire::TransportClosed { tool_id } => {
+                Self::NetworkError(format!("transport closed for {tool_id}"))
+            }
+            ToolErrorWire::UnsupportedProtocolVersion { supported } => {
+                Self::ProtocolError(format!("unsupported protocol; supported: {supported:?}"))
+            }
+            other => Self::Wire(other),
+        }
+    }
+}
+
+impl From<serde_json::Error> for ClientError {
+    fn from(err: serde_json::Error) -> Self {
+        Self::Serde(err.to_string())
+    }
+}
+
+impl From<IdError> for ClientError {
+    fn from(err: IdError) -> Self {
+        Self::ProtocolError(err.to_string())
+    }
+}
+
+impl From<url::ParseError> for ClientError {
+    fn from(err: url::ParseError) -> Self {
+        Self::InvalidConfig(format!("invalid url: {err}"))
+    }
+}
+
+impl From<tokio_tungstenite::tungstenite::Error> for ClientError {
+    fn from(err: tokio_tungstenite::tungstenite::Error) -> Self {
+        Self::NetworkError(err.to_string())
+    }
+}
+
+impl ClientError {
+    /// Classify a failed WebSocket upgrade. A `401`/`403` on the HTTP
+    /// upgrade is a non-retryable auth rejection
+    /// ([`Self::HandshakeAuthFailed`]); every other failure stays a
+    /// transport [`Self::NetworkError`] via the blanket `From` impl. The
+    /// distinction must be made here, before `From` collapses the typed
+    /// `Http` response status into an opaque string.
+    pub(crate) fn from_handshake_error(err: tokio_tungstenite::tungstenite::Error) -> Self {
+        if let tokio_tungstenite::tungstenite::Error::Http(resp) = &err {
+            let status = resp.status().as_u16();
+            if status == 401 || status == 403 {
+                return Self::HandshakeAuthFailed { status };
+            }
+        }
+        Self::from(err)
+    }
+}
+
+impl From<tokio::sync::oneshot::error::RecvError> for ClientError {
+    fn from(_: tokio::sync::oneshot::error::RecvError) -> Self {
+        Self::NetworkError("response waiter dropped (connection closed)".to_owned())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+    use xai_tool_protocol::{
+        WORKSPACE_UNAVAILABLE_SUBCODE, WorkspaceGonePhase, WorkspaceGoneReason,
+        workspace_unavailable_wire,
+    };
+
+    use super::*;
+
+    fn workspace_gone_envelope() -> JsonRpcError {
+        let wire = workspace_unavailable_wire(
+            WorkspaceGoneReason::Disconnect,
+            WorkspaceGonePhase::RouteMissing,
+        );
+        JsonRpcError {
+            code: -32005,
+            message: "workspace server gone".to_owned(),
+            data: Some(serde_json::to_value(&wire).unwrap()),
+        }
+    }
+
+    fn http_upgrade_error(status: u16) -> tokio_tungstenite::tungstenite::Error {
+        let resp = tokio_tungstenite::tungstenite::http::Response::builder()
+            .status(status)
+            .body(None::<Vec<u8>>)
+            .expect("response builds");
+        tokio_tungstenite::tungstenite::Error::Http(resp)
+    }
+
+    #[test]
+    fn handshake_401_and_403_map_to_handshake_auth_failed() {
+        for status in [401u16, 403] {
+            match ClientError::from_handshake_error(http_upgrade_error(status)) {
+                ClientError::HandshakeAuthFailed { status: got } => assert_eq!(got, status),
+                other => panic!("expected HandshakeAuthFailed for {status}; got {other:?}"),
+            }
+        }
+    }
+
+    #[test]
+    fn handshake_non_auth_status_stays_network_error() {
+        for status in [500u16, 502, 429] {
+            match ClientError::from_handshake_error(http_upgrade_error(status)) {
+                ClientError::NetworkError(_) => {}
+                other => panic!("expected NetworkError for {status}; got {other:?}"),
+            }
+        }
+    }
+
+    #[test]
+    fn from_jsonrpc_error_preserves_workspace_subcode_and_details() {
+        // The `data` payload decodes as `ToolErrorWire` first, so the stable
+        // subcode and structured details reach the SDK consumer intact rather
+        // than collapsing to the numeric code.
+        match ClientError::from_jsonrpc_error(workspace_gone_envelope()) {
+            ClientError::Wire(ToolErrorWire::Custom {
+                subcode, details, ..
+            }) => {
+                assert_eq!(subcode, WORKSPACE_UNAVAILABLE_SUBCODE);
+                let details = details.expect("details present");
+                assert_eq!(details["code"], json!(WORKSPACE_UNAVAILABLE_SUBCODE));
+                assert_eq!(details["reason"], json!("disconnect"));
+                assert_eq!(details["phase"], json!("route_missing"));
+                assert_eq!(details["retryable"], json!(true));
+            }
+            other => panic!("expected Wire(Custom), got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn is_server_not_found_recognizes_bare_minus_32601() {
+        // data-less -32601 -> custom subcode.
+        let err = ClientError::from_jsonrpc_error(JsonRpcError {
+            code: -32601,
+            message: "server abc not found for user".to_owned(),
+            data: None,
+        });
+        assert!(err.is_server_not_found());
+    }
+
+    #[test]
+    fn is_tool_unavailable_recognizes_bare_minus_32013() {
+        let err = ClientError::from_jsonrpc_error(JsonRpcError {
+            code: -32013,
+            message: "server abc did not complete the bind".to_owned(),
+            data: None,
+        });
+        assert!(err.is_tool_unavailable());
+    }
+
+    #[test]
+    fn is_server_not_found_rejects_other_errors() {
+        let auth = ClientError::from_jsonrpc_error(JsonRpcError {
+            code: -32002,
+            message: "nope".to_owned(),
+            data: None,
+        });
+        assert!(!auth.is_server_not_found());
+        // workspace-gone is the tool-call re-provision path, not bind ServerNotFound.
+        assert!(!ClientError::from_jsonrpc_error(workspace_gone_envelope()).is_server_not_found());
+    }
+
+    #[test]
+    fn bind_recognizers_are_mutually_exclusive() {
+        let not_found = ClientError::from_jsonrpc_error(JsonRpcError {
+            code: -32601,
+            message: "not found".to_owned(),
+            data: None,
+        });
+        let unavailable = ClientError::from_jsonrpc_error(JsonRpcError {
+            code: -32013,
+            message: "unavailable".to_owned(),
+            data: None,
+        });
+        assert!(not_found.is_server_not_found());
+        assert!(
+            !not_found.is_tool_unavailable(),
+            "-32601 must not be recognized as tool_unavailable"
+        );
+        assert!(unavailable.is_tool_unavailable());
+        assert!(
+            !unavailable.is_server_not_found(),
+            "-32013 must not be recognized as server_not_found"
+        );
+    }
+
+    #[test]
+    fn sdk_reexported_recognizer_matches_decoded_error() {
+        // SDK-only consumers reach the recognizer through the SDK re-export and
+        // the core decode path.
+        let err = xai_computer_hub_core::error_from_envelope(workspace_gone_envelope());
+        assert!(crate::is_workspace_unavailable(&err));
+    }
+
+    #[test]
+    fn sdk_reexported_recognizer_rejects_unrelated_custom_error() {
+        let wire = ToolErrorWire::Custom {
+            subcode: "unrelated".to_owned(),
+            message: "nope".to_owned(),
+            details: Some(json!({ "code": "unrelated" })),
+        };
+        let env = JsonRpcError {
+            code: -32000,
+            message: "nope".to_owned(),
+            data: Some(serde_json::to_value(&wire).unwrap()),
+        };
+        let err = xai_computer_hub_core::error_from_envelope(env);
+        assert!(!crate::is_workspace_unavailable(&err));
+    }
+}

crates/common/xai-computer-hub-sdk/src/handshake.rs

@@ -0,0 +1,93 @@
+//! Hello handshake helpers used by the connection actor and the
+//! reconnect-replay path.
+//!
+//! Splitting these into a dedicated module keeps the connection state
+//! machine readable: send the frame, parse the ack, surface a typed
+//! [`crate::ClientError`].
+
+use futures::{SinkExt, StreamExt};
+use tokio_tungstenite::tungstenite::Message;
+use xai_tool_protocol::{ConnectionKind, HelloAckMsg, HelloMsg};
+
+use crate::error::ClientError;
+
+/// Wire-protocol version both ends speak. Re-exported from the
+/// protocol crate so the SDK and the IC service share one source of
+/// truth.
+pub use xai_tool_protocol::PROTOCOL_VERSION;
+
+/// Send the [`HelloMsg`] and wait for the matching [`HelloAckMsg`].
+///
+/// `kind` should be [`ConnectionKind::ToolServer`] for tool-server
+/// builds (the only consumer today). The function returns the parsed
+/// ack so callers can observe the server-issued `connection_id` and
+/// the server-derived `user_id`.
+///
+/// When `server_id` is `Some`, it is included in the hello frame so the
+/// server can identify itself without a separate `register_server` call.
+pub async fn send_hello<Si, St>(
+    sink: &mut Si,
+    stream: &mut St,
+    kind: ConnectionKind,
+    server_id: Option<xai_tool_protocol::ServerId>,
+    description: Option<String>,
+    metadata: Option<serde_json::Value>,
+) -> Result<HelloAckMsg, ClientError>
+where
+    Si: SinkExt<Message> + Unpin,
+    Si::Error: std::fmt::Display,
+    St: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>> + Unpin,
+{
+    let hello = HelloMsg {
+        protocol_version: PROTOCOL_VERSION.to_owned(),
+        kind,
+        server_id,
+        description,
+        metadata,
+    };
+    let text = serde_json::to_string(&hello)?;
+    sink.send(Message::Text(text.into()))
+        .await
+        .map_err(|e| ClientError::NetworkError(format!("hello send failed: {e}")))?;
+
+    while let Some(msg) = stream.next().await {
+        let msg = msg?;
+        match msg {
+            Message::Text(text) => {
+                let ack: HelloAckMsg = serde_json::from_str(text.as_ref())
+                    .map_err(|e| ClientError::ProtocolError(format!("malformed hello_ack: {e}")))?;
+                if !ack
+                    .supported_protocol_versions
+                    .iter()
+                    .any(|v| v == PROTOCOL_VERSION)
+                {
+                    return Err(ClientError::ProtocolError(format!(
+                        "server does not support {PROTOCOL_VERSION}; supported: {:?}",
+                        ack.supported_protocol_versions
+                    )));
+                }
+                return Ok(ack);
+            }
+            Message::Ping(payload) => {
+                sink.send(Message::Pong(payload))
+                    .await
+                    .map_err(|e| ClientError::NetworkError(format!("pong send failed: {e}")))?;
+            }
+            Message::Close(frame) => {
+                let reason = frame.map(|f| f.reason.to_string()).unwrap_or_default();
+                return Err(ClientError::Closed(format!(
+                    "server closed during handshake: {reason}"
+                )));
+            }
+            Message::Pong(_) | Message::Frame(_) => continue,
+            Message::Binary(_) => {
+                return Err(ClientError::ProtocolError(
+                    "server sent binary frame during handshake".to_owned(),
+                ));
+            }
+        }
+    }
+    Err(ClientError::NetworkError(
+        "server closed before hello_ack".to_owned(),
+    ))
+}

crates/common/xai-computer-hub-sdk/src/lib.rs

@@ -0,0 +1,71 @@
+//! Tool-server and harness SDK.
+//!
+//! Single crate hosting both the tool-server runtime and the
+//! harness-side dispatch surface. The shared substrate —
+//! [`HubConnectionPool`], [`HubConnection`], the inbound demux, the
+//! refcount-managed bound-session set, and the transparent reconnect /
+//! replay state machine — lives here so both ends speak through one
+//! frame multiplex on top of one WebSocket per `(url, principal)`.
+//!
+//! The server entry point is [`ToolServer`]: build it via
+//! [`ToolServerBuilder`], wire one or more [`ToolServerHandler`]
+//! implementations, and call [`ToolServer::run`] to drive the inbound
+//! loop. The harness entry point is [`ToolHarness`]: build it via
+//! [`ToolHarnessBuilder`], optionally seed it with in-process
+//! [`xai_tool_runtime::Tool`] implementations, and call
+//! [`ToolHarness::call`] to dispatch a tool call. Authorisation
+//! credentials (`AuthCredential`) plus the target URL determine
+//! which pool entry the consumer attaches to; multiple
+//! [`ToolServer`] / [`ToolHarness`] instances against the same
+//! `(url, principal)` share a single connection and refcount their
+//! session bindings.
+
+#![forbid(unsafe_code)]
+
+pub(crate) mod admission;
+pub mod auth;
+pub(crate) mod cancel;
+pub mod connection;
+pub(crate) mod connection_borrow;
+pub mod demux;
+pub(crate) mod donate_pump;
+pub mod error;
+pub mod handshake;
+pub mod harness;
+pub mod log_donate;
+#[cfg(feature = "metrics")]
+pub mod metric_donate;
+pub mod metrics;
+pub mod notification;
+pub mod observability;
+pub mod pool;
+pub mod refcount;
+pub mod server;
+pub mod trace_donate;
+
+pub mod oidc_provider;
+
+pub use auth::{AuthCredential, AuthIdentity, AuthProvider, PrincipalKey, SharedAuthProvider};
+pub use connection::{ConnKey, HubConnection, ReconnectEvent};
+pub use error::ClientError;
+pub use harness::{
+    CancelOnDrop, LocalRegistry, ModelOutputExtractor, SessionBindReport, ToolHarness,
+    ToolHarnessBuilder, extractor_for,
+};
+pub use log_donate::{DonatingLogLayer, LogDonationPump, LogDonationSender, flush_log_layer};
+#[cfg(feature = "metrics")]
+pub use metric_donate::MetricDonationPump;
+pub use notification::HubNotification;
+pub use observability::ObservabilityBridge;
+pub use oidc_provider::{
+    OidcAuthProvider, OidcAuthProviderBuilder, OnRefreshCallback, RefreshEvent,
+};
+pub use pool::HubConnectionPool;
+pub use server::{
+    ResolvedSessionHandlers, SessionHandlerResolver, SystemNotifyAck, ToolServer,
+    ToolServerBuilder, ToolServerHandler, WeakToolServer,
+};
+pub use trace_donate::{HubDonatingReporter, TraceDonationPump};
+// Re-exported so consumers that depend only on the SDK can recognize the
+// server's `workspace_unavailable` error without also pulling in the core crate.
+pub use xai_computer_hub_core::is_workspace_unavailable;

crates/common/xai-computer-hub-sdk/src/log_donate.rs

@@ -0,0 +1,592 @@
+//! Forward curated `tracing` events to the connected server over the
+//! WebSocket transport (`logs.donate`).
+//!
+//! [`DonatingLogLayer`] is installed **inert** at startup and activated
+//! post-connect by swapping in a [`LogDonationSender`] (a global-subscriber
+//! constraint); while inert, selected events are dropped before enqueueing.
+//!
+//! Only events on the [`TELEMETRY_TARGET`] target at `>= INFO` are
+//! forwarded, and only fields in [`ALLOWED_FIELDS`] are included; other
+//! fields such as `reason`/`error` are omitted.
+
+use std::sync::{Arc, LazyLock};
+use std::time::{Duration, Instant};
+
+use arc_swap::ArcSwapOption;
+use base64::Engine as _;
+use fastrace::collector::SpanContext;
+use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest;
+use opentelemetry_proto::tonic::common::v1::{AnyValue, InstrumentationScope, KeyValue, any_value};
+use opentelemetry_proto::tonic::logs::v1::{LogRecord, ResourceLogs, ScopeLogs};
+use opentelemetry_proto::tonic::resource::v1::Resource;
+use prost::Message as _;
+use tokio::sync::mpsc;
+use tracing::Level;
+use tracing::field::{Field, Visit};
+use tracing_subscriber::Layer;
+use tracing_subscriber::layer::Context;
+use xai_tool_protocol::{MAX_DONATION_BYTES, MAX_LOG_RECORDS_PER_DONATION};
+
+use crate::donate_pump::{
+    PENDING_FLUSHES, PumpMsg, drain_via, make_resource, now_unix_nanos, run_pump, string_kv,
+    string_value,
+};
+use crate::server::ToolServer;
+
+/// Stable target the workspace routes selected events through.
+/// The layer selects exactly this target, ignoring global
+/// `RUST_LOG`. The server re-stamps it as the OTLP scope name.
+pub const TELEMETRY_TARGET: &str = "workspace::telemetry";
+
+/// Set of forwardable field names — guaranteed-literal or numeric.
+/// Only the listed fields are included; other fields such as
+/// `reason`/`error`/`object_path`/`gcs_path` are omitted.
+const ALLOWED_FIELDS: &[&str] = &[
+    "session_id",
+    "turn_number",
+    "phase",
+    "bytes",
+    "file_count",
+    "pending",
+    "pending_bytes",
+    "sample_period_secs",
+    "error_category",
+    "outcome",
+    "skip_reason",
+    "drain_reason",
+    "grace_ms",
+    "active_at_start",
+    "pending_at_start",
+    "producers_at_start",
+];
+
+/// Flush a buffered batch once it reaches this many records.
+const LOG_BATCH_FLUSH_RECORDS: usize = 32;
+/// Flush a partial batch once its oldest record is at least this old
+/// (checked on the next event; the tail is fenced by teardown).
+const LOG_BATCH_MAX_AGE: Duration = Duration::from_secs(2);
+
+fn is_allowed(name: &str) -> bool {
+    ALLOWED_FIELDS.contains(&name)
+}
+
+/// `tracing::Level` → OTLP (`SeverityText`, `SeverityNumber`).
+fn severity(level: &Level) -> (&'static str, i32) {
+    match *level {
+        Level::ERROR => ("ERROR", 17),
+        Level::WARN => ("WARN", 13),
+        Level::INFO => ("INFO", 9),
+        Level::DEBUG => ("DEBUG", 5),
+        Level::TRACE => ("TRACE", 1),
+    }
+}
+
+/// `>= INFO` in severity terms (INFO/WARN/ERROR). Note tracing orders
+/// `ERROR < WARN < INFO < DEBUG < TRACE`, so this is `level <= INFO`.
+fn at_least_info(level: &Level) -> bool {
+    *level <= Level::INFO
+}
+
+/// Big-endian byte encoding of the local parent's ids into the OTLP
+/// 16-byte / 8-byte fields; empty when no fastrace local parent is
+/// active (the common case for detached producer tasks).
+fn current_ids() -> (Vec<u8>, Vec<u8>) {
+    match SpanContext::current_local_parent() {
+        Some(ctx) => encode_ids(&ctx),
+        None => (Vec::new(), Vec::new()),
+    }
+}
+
+fn encode_ids(ctx: &SpanContext) -> (Vec<u8>, Vec<u8>) {
+    (
+        ctx.trace_id.0.to_be_bytes().to_vec(),
+        ctx.span_id.0.to_be_bytes().to_vec(),
+    )
+}
+
+/// Field visitor: keeps the message as the OTLP `Body` and only
+/// allowlisted fields as attributes; everything else is dropped.
+#[derive(Default)]
+struct AllowlistVisitor {
+    body: Option<String>,
+    attributes: Vec<KeyValue>,
+}
+
+impl Visit for AllowlistVisitor {
+    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
+        let name = field.name();
+        if name == "message" {
+            self.body = Some(format!("{value:?}"));
+        } else if is_allowed(name) {
+            self.attributes.push(string_kv(name, format!("{value:?}")));
+        }
+    }
+
+    fn record_str(&mut self, field: &Field, value: &str) {
+        let name = field.name();
+        if name == "message" {
+            self.body = Some(value.to_owned());
+        } else if is_allowed(name) {
+            self.attributes.push(string_kv(name, value.to_owned()));
+        }
+    }
+
+    fn record_i64(&mut self, field: &Field, value: i64) {
+        if is_allowed(field.name()) {
+            self.push_int(field.name(), value);
+        }
+    }
+
+    fn record_u64(&mut self, field: &Field, value: u64) {
+        if is_allowed(field.name()) {
+            self.push_int(field.name(), value as i64);
+        }
+    }
+
+    fn record_bool(&mut self, field: &Field, value: bool) {
+        if is_allowed(field.name()) {
+            self.attributes.push(KeyValue {
+                key: field.name().to_owned(),
+                value: Some(AnyValue {
+                    value: Some(any_value::Value::BoolValue(value)),
+                }),
+                ..Default::default()
+            });
+        }
+    }
+
+    fn record_f64(&mut self, field: &Field, value: f64) {
+        if is_allowed(field.name()) {
+            self.attributes.push(KeyValue {
+                key: field.name().to_owned(),
+                value: Some(AnyValue {
+                    value: Some(any_value::Value::DoubleValue(value)),
+                }),
+                ..Default::default()
+            });
+        }
+    }
+}
+
+impl AllowlistVisitor {
+    fn push_int(&mut self, name: &str, value: i64) {
+        self.attributes.push(KeyValue {
+            key: name.to_owned(),
+            value: Some(AnyValue {
+                value: Some(any_value::Value::IntValue(value)),
+            }),
+            ..Default::default()
+        });
+    }
+}
+
+fn build_log_record(level: &Level, visitor: AllowlistVisitor) -> LogRecord {
+    let now_nanos = now_unix_nanos();
+    let (text, number) = severity(level);
+    let (trace_id, span_id) = current_ids();
+    LogRecord {
+        time_unix_nano: now_nanos,
+        observed_time_unix_nano: now_nanos,
+        severity_number: number,
+        severity_text: text.to_owned(),
+        body: visitor.body.map(string_value),
+        attributes: visitor.attributes,
+        trace_id,
+        span_id,
+        ..Default::default()
+    }
+}
+
+/// Encodes batches of OTLP `LogRecord`s onto the pump channel. Chunks at
+/// [`MAX_LOG_RECORDS_PER_DONATION`], drops payloads over
+/// [`MAX_DONATION_BYTES`], and never blocks.
+#[derive(Clone)]
+struct PumpLogExporter {
+    tx: mpsc::Sender<PumpMsg>,
+    resource: Resource,
+}
+
+impl PumpLogExporter {
+    fn export(&self, mut records: Vec<LogRecord>) {
+        while !records.is_empty() {
+            let chunk = if records.len() > MAX_LOG_RECORDS_PER_DONATION {
+                let rest = records.split_off(MAX_LOG_RECORDS_PER_DONATION);
+                std::mem::replace(&mut records, rest)
+            } else {
+                std::mem::take(&mut records)
+            };
+            let request = ExportLogsServiceRequest {
+                resource_logs: vec![ResourceLogs {
+                    resource: Some(self.resource.clone()),
+                    scope_logs: vec![ScopeLogs {
+                        scope: Some(InstrumentationScope {
+                            name: TELEMETRY_TARGET.to_owned(),
+                            ..Default::default()
+                        }),
+                        log_records: chunk,
+                        schema_url: String::new(),
+                    }],
+                    schema_url: String::new(),
+                }],
+            };
+            let bytes = request.encode_to_vec();
+            if bytes.len() > MAX_DONATION_BYTES {
+                tracing::debug!(len = bytes.len(), "dropping oversized log donation payload");
+                continue;
+            }
+            let payload = base64::engine::general_purpose::STANDARD.encode(bytes);
+            if self.tx.try_send(PumpMsg::Payload(payload)).is_err() {
+                tracing::debug!("log donation queue full; dropping log batch");
+            }
+        }
+    }
+}
+
+/// Activation handle swapped into an inert [`DonatingLogLayer`]. Wraps
+/// the pump sender plus the resource (`service.name`) the layer needs to
+/// encode batches.
+pub struct LogDonationSender {
+    exporter: PumpLogExporter,
+}
+
+impl LogDonationSender {
+    fn export(&self, records: Vec<LogRecord>) {
+        self.exporter.export(records);
+    }
+}
+
+#[derive(Default)]
+struct LogBatch {
+    records: Vec<LogRecord>,
+    oldest: Option<Instant>,
+}
+
+struct LogLayerShared {
+    sender: ArcSwapOption<LogDonationSender>,
+    batch: parking_lot::Mutex<LogBatch>,
+}
+
+impl LogLayerShared {
+    /// Buffer a record; return any records due for flush (count/age).
+    fn push(&self, record: LogRecord) -> Vec<LogRecord> {
+        let mut batch = self.batch.lock();
+        if batch.records.is_empty() {
+            batch.oldest = Some(Instant::now());
+        }
+        batch.records.push(record);
+        let due = batch.records.len() >= LOG_BATCH_FLUSH_RECORDS
+            || batch
+                .oldest
+                .is_some_and(|t| t.elapsed() >= LOG_BATCH_MAX_AGE);
+        if due {
+            batch.oldest = None;
+            std::mem::take(&mut batch.records)
+        } else {
+            Vec::new()
+        }
+    }
+
+    /// Force the buffered batch onto the pump (teardown analogue of
+    /// `fastrace::flush()`); no-op while inert or empty.
+    fn flush(&self) {
+        let records = {
+            let mut batch = self.batch.lock();
+            batch.oldest = None;
+            std::mem::take(&mut batch.records)
+        };
+        if records.is_empty() {
+            return;
+        }
+        if let Some(sender) = self.sender.load_full() {
+            sender.export(records);
+        }
+    }
+}
+
+/// Process-global handle to the active layer's shared state so
+/// [`flush_log_layer`] can drive a teardown flush without a reference.
+static ACTIVE_LOG_LAYER: LazyLock<ArcSwapOption<LogLayerShared>> =
+    LazyLock::new(ArcSwapOption::empty);
+
+/// A composable [`tracing_subscriber::Layer`] that converts selected
+/// events into OTLP log records and batches them onto the pump.
+/// Installed inert; activated by [`Self::activate`].
+#[derive(Clone)]
+pub struct DonatingLogLayer {
+    shared: Arc<LogLayerShared>,
+}
+
+impl DonatingLogLayer {
+    /// Install inert (no sender): selected events are dropped until
+    /// [`Self::activate`] swaps a sender in. Registers itself as the
+    /// process-global flush target.
+    pub fn new_inert() -> Self {
+        let shared = Arc::new(LogLayerShared {
+            sender: ArcSwapOption::empty(),
+            batch: parking_lot::Mutex::new(LogBatch::default()),
+        });
+        ACTIVE_LOG_LAYER.store(Some(shared.clone()));
+        Self { shared }
+    }
+
+    /// Swap in the donation sender, activating donation.
+    pub fn activate(&self, sender: LogDonationSender) {
+        self.shared.sender.store(Some(Arc::new(sender)));
+    }
+}
+
+impl<S: tracing::Subscriber> Layer<S> for DonatingLogLayer {
+    fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
+        let Some(sender) = self.shared.sender.load_full() else {
+            return;
+        };
+        let meta = event.metadata();
+        if meta.target() != TELEMETRY_TARGET || !at_least_info(meta.level()) {
+            return;
+        }
+        let mut visitor = AllowlistVisitor::default();
+        event.record(&mut visitor);
+        let record = build_log_record(meta.level(), visitor);
+        let due = self.shared.push(record);
+        if !due.is_empty() {
+            sender.export(due);
+        }
+    }
+}
+
+/// Flush the active [`DonatingLogLayer`]'s in-memory batch onto the
+/// pump. Called from `ToolServer` teardown before the pump drain so a
+/// crash-y shutdown does not abandon a partial batch.
+pub fn flush_log_layer() {
+    if let Some(shared) = ACTIVE_LOG_LAYER.load_full() {
+        shared.flush();
+    }
+}
+
+/// Shutdown fence: drains queued log donations before the connection
+/// closes. Call after [`flush_log_layer`].
+pub struct LogDonationPump {
+    tx: mpsc::Sender<PumpMsg>,
+}
+
+impl LogDonationPump {
+    /// Resolves once every payload queued before this call has had a
+    /// send attempt.
+    pub async fn drain(&self) {
+        drain_via(&self.tx).await;
+    }
+}
+
+impl ToolServer {
+    /// Post-connect entry point: spawn the log donation pump (wiring
+    /// [`ToolServer::donate_logs`]) and return a sender to swap into the
+    /// already-installed inert [`DonatingLogLayer`] plus a drain handle.
+    /// Does **not** return a `Layer` — a layer cannot be added to an
+    /// already-set global subscriber. `service_name` must be
+    /// server-allowlisted.
+    pub fn log_donation_layer(
+        &self,
+        service_name: impl Into<String>,
+    ) -> (LogDonationSender, LogDonationPump) {
+        let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
+        let server = self.downgrade();
+        tokio::spawn(run_pump(rx, move |payload: String| {
+            let server = server.clone();
+            async move {
+                let Some(server) = server.upgrade() else {
+                    return (false, payload);
+                };
+                let ok = server.donate_logs(&payload).await.is_ok();
+                (ok, payload)
+            }
+        }));
+        self.set_log_donation_pump(tx.clone());
+
+        let sender = LogDonationSender {
+            exporter: PumpLogExporter {
+                tx: tx.clone(),
+                resource: make_resource(service_name.into()),
+            },
+        };
+        (sender, LogDonationPump { tx })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use fastrace::collector::{SpanId, TraceId};
+    use tracing_subscriber::layer::SubscriberExt;
+
+    use super::*;
+
+    fn test_sender(tx: mpsc::Sender<PumpMsg>) -> LogDonationSender {
+        LogDonationSender {
+            exporter: PumpLogExporter {
+                tx,
+                resource: make_resource("test-service".to_owned()),
+            },
+        }
+    }
+
+    fn decode(payload: String) -> ExportLogsServiceRequest {
+        let bytes = base64::engine::general_purpose::STANDARD
+            .decode(payload)
+            .expect("payload must be base64");
+        ExportLogsServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP")
+    }
+
+    #[test]
+    fn severity_maps_levels_to_otlp_numbers() {
+        assert_eq!(severity(&Level::ERROR), ("ERROR", 17));
+        assert_eq!(severity(&Level::WARN), ("WARN", 13));
+        assert_eq!(severity(&Level::INFO), ("INFO", 9));
+        assert_eq!(severity(&Level::DEBUG), ("DEBUG", 5));
+        assert_eq!(severity(&Level::TRACE), ("TRACE", 1));
+    }
+
+    #[test]
+    fn donation_filter_selects_info_and_above() {
+        assert!(at_least_info(&Level::ERROR));
+        assert!(at_least_info(&Level::WARN));
+        assert!(at_least_info(&Level::INFO));
+        assert!(!at_least_info(&Level::DEBUG));
+        assert!(!at_least_info(&Level::TRACE));
+    }
+
+    #[test]
+    fn encode_ids_is_big_endian_16_and_8_bytes() {
+        let ctx = SpanContext::new(
+            TraceId(0x0af7651916cd43dd8448eb211c80319c),
+            SpanId(0xb7ad6b7169203331),
+        );
+        let (trace_id, span_id) = encode_ids(&ctx);
+        assert_eq!(trace_id.len(), 16);
+        assert_eq!(span_id.len(), 8);
+        assert_eq!(
+            format!("{:032x}", u128::from_be_bytes(trace_id.try_into().unwrap())),
+            "0af7651916cd43dd8448eb211c80319c"
+        );
+        assert_eq!(
+            format!("{:016x}", u64::from_be_bytes(span_id.try_into().unwrap())),
+            "b7ad6b7169203331"
+        );
+    }
+
+    #[test]
+    fn current_ids_empty_without_local_parent() {
+        let (trace_id, span_id) = current_ids();
+        assert!(trace_id.is_empty());
+        assert!(span_id.is_empty());
+    }
+
+    #[test]
+    fn layer_converts_event_and_redacts_free_form_fields() {
+        let layer = DonatingLogLayer::new_inert();
+        let (tx, mut rx) = mpsc::channel::<PumpMsg>(8);
+        layer.activate(test_sender(tx));
+        let flusher = layer.clone();
+
+        let subscriber = tracing_subscriber::registry().with(layer);
+        tracing::subscriber::with_default(subscriber, || {
+            tracing::warn!(
+                target: "workspace::telemetry",
+                session_id = "s1",
+                turn_number = 3u64,
+                phase = "tool_state",
+                error_category = "archive_failed",
+                error = "secret git stderr with /home/user/path",
+                "archive build failed (queued path)"
+            );
+            // Off-target event must never be forwarded.
+            tracing::warn!(session_id = "s2", "unrelated chatter");
+            // DEBUG on-target is below the threshold.
+            tracing::debug!(target: "workspace::telemetry", session_id = "s3", "verbose");
+        });
+        flusher.shared.flush();
+
+        let PumpMsg::Payload(payload) = rx.try_recv().expect("one batch must be queued") else {
+            panic!("expected a payload");
+        };
+        let request = decode(payload);
+        let scope_logs = &request.resource_logs[0].scope_logs[0];
+        assert_eq!(
+            scope_logs.scope.as_ref().unwrap().name,
+            "workspace::telemetry"
+        );
+        assert_eq!(
+            scope_logs.log_records.len(),
+            1,
+            "only the WARN on-target row"
+        );
+
+        let record = &scope_logs.log_records[0];
+        assert_eq!(record.severity_text, "WARN");
+        assert_eq!(record.severity_number, 13);
+        assert_eq!(
+            record.body.as_ref().unwrap().value,
+            Some(any_value::Value::StringValue(
+                "archive build failed (queued path)".to_owned()
+            ))
+        );
+        let keys: Vec<&str> = record.attributes.iter().map(|kv| kv.key.as_str()).collect();
+        assert!(keys.contains(&"session_id"));
+        assert!(keys.contains(&"turn_number"));
+        assert!(keys.contains(&"phase"));
+        assert!(keys.contains(&"error_category"));
+        assert!(
+            !keys.contains(&"error"),
+            "free-form `error` must be dropped, got {keys:?}"
+        );
+
+        // Resource carries the donor service.name.
+        let service_name = request.resource_logs[0]
+            .resource
+            .as_ref()
+            .unwrap()
+            .attributes
+            .iter()
+            .find(|kv| kv.key == "service.name")
+            .and_then(|kv| kv.value.as_ref())
+            .and_then(|v| v.value.clone());
+        assert_eq!(
+            service_name,
+            Some(any_value::Value::StringValue("test-service".to_owned()))
+        );
+
+        assert!(rx.try_recv().is_err(), "no further payloads");
+    }
+
+    #[test]
+    fn inert_layer_drops_selected_events() {
+        let layer = DonatingLogLayer::new_inert();
+        let flusher = layer.clone();
+        // No sender activated.
+        let subscriber = tracing_subscriber::registry().with(layer);
+        tracing::subscriber::with_default(subscriber, || {
+            tracing::warn!(target: "workspace::telemetry", session_id = "s1", "dropped");
+        });
+        flusher.shared.flush();
+        // Nothing to assert beyond not panicking: with no sender the
+        // batch never fills and flush is a no-op.
+    }
+
+    #[test]
+    fn exporter_chunks_at_max_records_per_donation() {
+        let (tx, mut rx) = mpsc::channel::<PumpMsg>(8);
+        let exporter = PumpLogExporter {
+            tx,
+            resource: make_resource("test-service".to_owned()),
+        };
+        let records = vec![LogRecord::default(); MAX_LOG_RECORDS_PER_DONATION + 1];
+        exporter.export(records);
+
+        let mut total = 0;
+        let mut payloads = 0;
+        while let Ok(PumpMsg::Payload(p)) = rx.try_recv() {
+            payloads += 1;
+            total += decode(p).resource_logs[0].scope_logs[0].log_records.len();
+        }
+        assert_eq!(payloads, 2, "one full chunk + remainder");
+        assert_eq!(total, MAX_LOG_RECORDS_PER_DONATION + 1);
+    }
+}

crates/common/xai-computer-hub-sdk/src/metric_donate.rs

@@ -0,0 +1,396 @@
+//! Forward the process's Prometheus metrics to the connected server over
+//! the WebSocket transport (`metrics.donate`).
+//!
+//! A [`MetricDonationReporter`] periodically snapshots the default
+//! Prometheus registry via [`prometheus::gather`], converts the
+//! `MetricFamily` set to native OTLP metrics (Counter→Sum, Gauge→Gauge,
+//! Histogram→Histogram, labels preserved, cumulative temporality), and
+//! pumps the batch over the shared [`crate::donate_pump`]. Because it
+//! gathers the whole registry, every current and future metric is
+//! exported with zero per-metric wiring. Metrics are **process-aggregate**
+//! — [`ToolServer::donate_metrics`] requires no bound session.
+
+use std::sync::{Arc, LazyLock};
+use std::time::Duration;
+
+use arc_swap::ArcSwapOption;
+use base64::Engine as _;
+use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest;
+use opentelemetry_proto::tonic::common::v1::KeyValue;
+use opentelemetry_proto::tonic::metrics::v1::{
+    AggregationTemporality, Gauge, Histogram, HistogramDataPoint, Metric, NumberDataPoint,
+    ResourceMetrics, ScopeMetrics, Sum, metric, number_data_point,
+};
+use opentelemetry_proto::tonic::resource::v1::Resource;
+use prometheus::proto::{MetricFamily, MetricType};
+use prost::Message as _;
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+use xai_tool_protocol::{MAX_DONATION_BYTES, MAX_METRICS_PER_DONATION};
+
+use crate::donate_pump::{
+    PENDING_FLUSHES, PumpMsg, make_resource, now_unix_nanos, run_pump, string_kv,
+};
+use crate::server::ToolServer;
+
+/// How often the reporter snapshots the registry. The server re-stamps
+/// attribution; cumulative temporality means missed ticks only delay
+/// freshness, never lose monotonic state.
+const DEFAULT_GATHER_INTERVAL: Duration = Duration::from_secs(60);
+
+fn labels_to_kv(labels: &[prometheus::proto::LabelPair]) -> Vec<KeyValue> {
+    labels
+        .iter()
+        .map(|l| string_kv(l.name(), l.value().to_owned()))
+        .collect()
+}
+
+fn number_point(metric: &prometheus::proto::Metric, value: f64, now: u64) -> NumberDataPoint {
+    NumberDataPoint {
+        attributes: labels_to_kv(metric.get_label()),
+        time_unix_nano: now,
+        value: Some(number_data_point::Value::AsDouble(value)),
+        ..Default::default()
+    }
+}
+
+/// Prometheus histogram buckets are **cumulative** (`le` counts); OTLP
+/// wants per-bucket counts plus an implicit `+Inf` bucket, so the
+/// cumulative counts are differenced here.
+fn histogram_point(metric: &prometheus::proto::Metric, now: u64) -> HistogramDataPoint {
+    let hist = metric.get_histogram();
+    let mut bucket_counts = Vec::new();
+    let mut explicit_bounds = Vec::new();
+    let mut prev = 0u64;
+    for bucket in hist.get_bucket() {
+        let cumulative = bucket.cumulative_count();
+        bucket_counts.push(cumulative.saturating_sub(prev));
+        explicit_bounds.push(bucket.upper_bound());
+        prev = cumulative;
+    }
+    let total = hist.get_sample_count();
+    bucket_counts.push(total.saturating_sub(prev));
+    HistogramDataPoint {
+        attributes: labels_to_kv(metric.get_label()),
+        time_unix_nano: now,
+        count: total,
+        sum: Some(hist.get_sample_sum()),
+        bucket_counts,
+        explicit_bounds,
+        ..Default::default()
+    }
+}
+
+/// Convert a gathered `MetricFamily` set to OTLP metrics. Summaries and
+/// untyped families are skipped defensively (none registered today).
+fn convert_families(families: &[MetricFamily]) -> Vec<Metric> {
+    let now = now_unix_nanos();
+    let cumulative = AggregationTemporality::Cumulative as i32;
+    let mut out = Vec::new();
+    for family in families {
+        let name = family.name().to_owned();
+        let data = match family.get_field_type() {
+            MetricType::COUNTER => metric::Data::Sum(Sum {
+                data_points: family
+                    .get_metric()
+                    .iter()
+                    .map(|m| number_point(m, m.get_counter().value(), now))
+                    .collect(),
+                aggregation_temporality: cumulative,
+                is_monotonic: true,
+            }),
+            MetricType::GAUGE => metric::Data::Gauge(Gauge {
+                data_points: family
+                    .get_metric()
+                    .iter()
+                    .map(|m| number_point(m, m.get_gauge().value(), now))
+                    .collect(),
+            }),
+            MetricType::HISTOGRAM => metric::Data::Histogram(Histogram {
+                data_points: family
+                    .get_metric()
+                    .iter()
+                    .map(|m| histogram_point(m, now))
+                    .collect(),
+                aggregation_temporality: cumulative,
+            }),
+            MetricType::SUMMARY | MetricType::UNTYPED => continue,
+        };
+        out.push(Metric {
+            name,
+            data: Some(data),
+            ..Default::default()
+        });
+    }
+    out
+}
+
+/// Encodes batches of OTLP metrics onto the pump channel. Chunks at
+/// [`MAX_METRICS_PER_DONATION`], drops payloads over
+/// [`MAX_DONATION_BYTES`], and never blocks.
+#[derive(Clone)]
+struct MetricExporter {
+    tx: mpsc::Sender<PumpMsg>,
+    resource: Resource,
+}
+
+impl MetricExporter {
+    fn export(&self, mut metrics: Vec<Metric>) {
+        while !metrics.is_empty() {
+            let chunk = if metrics.len() > MAX_METRICS_PER_DONATION {
+                let rest = metrics.split_off(MAX_METRICS_PER_DONATION);
+                std::mem::replace(&mut metrics, rest)
+            } else {
+                std::mem::take(&mut metrics)
+            };
+            let request = ExportMetricsServiceRequest {
+                resource_metrics: vec![ResourceMetrics {
+                    resource: Some(self.resource.clone()),
+                    scope_metrics: vec![ScopeMetrics {
+                        metrics: chunk,
+                        ..Default::default()
+                    }],
+                    schema_url: String::new(),
+                }],
+            };
+            let bytes = request.encode_to_vec();
+            if bytes.len() > MAX_DONATION_BYTES {
+                tracing::debug!(
+                    len = bytes.len(),
+                    "dropping oversized metric donation payload"
+                );
+                continue;
+            }
+            let payload = base64::engine::general_purpose::STANDARD.encode(bytes);
+            if self.tx.try_send(PumpMsg::Payload(payload)).is_err() {
+                tracing::debug!("metric donation queue full; dropping metric batch");
+            }
+        }
+    }
+
+    fn gather_and_send(&self) {
+        let metrics = convert_families(&prometheus::gather());
+        if metrics.is_empty() {
+            return;
+        }
+        self.export(metrics);
+    }
+}
+
+/// Process-global handle to the active exporter so [`gather_and_send`]
+/// can drive a final teardown gather without a reference.
+static ACTIVE_METRIC_EXPORTER: LazyLock<ArcSwapOption<MetricExporter>> =
+    LazyLock::new(ArcSwapOption::empty);
+
+/// Final registry gather onto the active metric pump. Called from
+/// `ToolServer` teardown before the pump drain so a crash-y shutdown
+/// captures the latest values.
+pub(crate) fn gather_and_send() {
+    if let Some(exporter) = ACTIVE_METRIC_EXPORTER.load_full() {
+        exporter.gather_and_send();
+    }
+}
+
+/// Drop the process-global exporter on teardown so its pump `Sender` is released
+/// and the metric pump can wind down. Called from `flush_donations_inner` after
+/// the final [`gather_and_send`] (and alongside clearing the stored pump
+/// senders), so a dropped `ToolServer` doesn't leak the pump task.
+pub(crate) fn clear_active_exporter() {
+    ACTIVE_METRIC_EXPORTER.store(None);
+}
+
+/// Periodic registry gatherer spawned by
+/// [`ToolServer::metric_donation_reporter`]. Internal: constructed and run
+/// only by `metric_donation_reporter`; not part of the crate's public API.
+pub(crate) struct MetricDonationReporter {
+    exporter: MetricExporter,
+    interval: Duration,
+    shutdown: CancellationToken,
+}
+
+impl MetricDonationReporter {
+    async fn run(self) {
+        let mut ticker = tokio::time::interval(self.interval);
+        loop {
+            tokio::select! {
+                _ = ticker.tick() => self.exporter.gather_and_send(),
+                // Stop on teardown so this task (and the pump `tx` clone it
+                // holds) doesn't outlive `ToolServer::shutdown` and keep
+                // gathering/sending forever.
+                _ = self.shutdown.cancelled() => break,
+            }
+        }
+    }
+}
+
+/// Shutdown fence: drains queued metric donations before the connection
+/// closes.
+pub struct MetricDonationPump {
+    tx: mpsc::Sender<PumpMsg>,
+}
+
+impl MetricDonationPump {
+    /// Resolves once every payload queued before this call has had a
+    /// send attempt.
+    pub async fn drain(&self) {
+        crate::donate_pump::drain_via(&self.tx).await;
+    }
+}
+
+impl ToolServer {
+    /// Post-connect entry point: spawn the metric donation pump (wiring
+    /// [`ToolServer::donate_metrics`]) plus the periodic registry
+    /// gatherer, and return a drain handle. Activates on server presence
+    /// (no env flag). `service_name` must be server-allowlisted.
+    pub fn metric_donation_reporter(&self, service_name: impl Into<String>) -> MetricDonationPump {
+        let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
+        let server = self.downgrade();
+        tokio::spawn(run_pump(rx, move |payload: String| {
+            let server = server.clone();
+            async move {
+                let Some(server) = server.upgrade() else {
+                    return (false, payload);
+                };
+                let ok = server.donate_metrics(&payload).await.is_ok();
+                (ok, payload)
+            }
+        }));
+        self.set_metric_donation_pump(tx.clone());
+
+        let exporter = MetricExporter {
+            tx: tx.clone(),
+            resource: make_resource(service_name.into()),
+        };
+        ACTIVE_METRIC_EXPORTER.store(Some(Arc::new(exporter.clone())));
+        tokio::spawn(
+            MetricDonationReporter {
+                exporter,
+                interval: DEFAULT_GATHER_INTERVAL,
+                shutdown: self.shutdown_token(),
+            }
+            .run(),
+        );
+
+        MetricDonationPump { tx }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use opentelemetry_proto::tonic::common::v1::any_value;
+    use prometheus::{Histogram, HistogramOpts, IntCounterVec, IntGauge, Opts, Registry};
+
+    use super::*;
+
+    fn decode(payload: String) -> ExportMetricsServiceRequest {
+        let bytes = base64::engine::general_purpose::STANDARD
+            .decode(payload)
+            .expect("payload must be base64");
+        ExportMetricsServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP")
+    }
+
+    fn label_map(attrs: &[KeyValue]) -> std::collections::HashMap<String, String> {
+        attrs
+            .iter()
+            .filter_map(|kv| match kv.value.as_ref().and_then(|v| v.value.clone()) {
+                Some(any_value::Value::StringValue(s)) => Some((kv.key.clone(), s)),
+                _ => None,
+            })
+            .collect()
+    }
+
+    #[test]
+    fn converts_counter_gauge_histogram_with_labels() {
+        let registry = Registry::new();
+
+        let counter =
+            IntCounterVec::new(Opts::new("grok_test_total", "help"), &["reason"]).unwrap();
+        registry.register(Box::new(counter.clone())).unwrap();
+        counter.with_label_values(&["zdr"]).inc_by(5);
+
+        let gauge = IntGauge::new("grok_test_pending", "help").unwrap();
+        registry.register(Box::new(gauge.clone())).unwrap();
+        gauge.set(7);
+
+        let hist = Histogram::with_opts(
+            HistogramOpts::new("grok_test_seconds", "help").buckets(vec![0.5, 1.0]),
+        )
+        .unwrap();
+        registry.register(Box::new(hist.clone())).unwrap();
+        hist.observe(0.25);
+        hist.observe(0.75);
+        hist.observe(5.0);
+
+        let metrics = convert_families(&registry.gather());
+        let by_name: std::collections::HashMap<_, _> =
+            metrics.iter().map(|m| (m.name.clone(), m)).collect();
+
+        // Counter -> Sum (monotonic, cumulative), label preserved.
+        let metric::Data::Sum(sum) = by_name["grok_test_total"].data.as_ref().unwrap() else {
+            panic!("counter must convert to Sum");
+        };
+        assert!(sum.is_monotonic);
+        assert_eq!(
+            sum.aggregation_temporality,
+            AggregationTemporality::Cumulative as i32
+        );
+        let dp = &sum.data_points[0];
+        assert_eq!(dp.value, Some(number_data_point::Value::AsDouble(5.0)));
+        assert_eq!(
+            label_map(&dp.attributes).get("reason").map(String::as_str),
+            Some("zdr")
+        );
+
+        // Gauge -> Gauge.
+        let metric::Data::Gauge(g) = by_name["grok_test_pending"].data.as_ref().unwrap() else {
+            panic!("gauge must convert to Gauge");
+        };
+        assert_eq!(
+            g.data_points[0].value,
+            Some(number_data_point::Value::AsDouble(7.0))
+        );
+
+        // Histogram -> Histogram with cumulative buckets differenced and
+        // a +Inf bucket appended.
+        let metric::Data::Histogram(h) = by_name["grok_test_seconds"].data.as_ref().unwrap() else {
+            panic!("histogram must convert to Histogram");
+        };
+        assert_eq!(
+            h.aggregation_temporality,
+            AggregationTemporality::Cumulative as i32
+        );
+        let hdp = &h.data_points[0];
+        assert_eq!(hdp.count, 3);
+        assert_eq!(hdp.sum, Some(6.0));
+        assert_eq!(hdp.explicit_bounds, vec![0.5, 1.0]);
+        // (<=0.5): 0.25 -> 1 ; (0.5,1.0]: 0.75 -> 1 ; (+Inf): 5.0 -> 1
+        assert_eq!(hdp.bucket_counts, vec![1, 1, 1]);
+    }
+
+    #[test]
+    fn exporter_chunks_at_max_metrics_per_donation() {
+        let (tx, mut rx) = mpsc::channel::<PumpMsg>(8);
+        let exporter = MetricExporter {
+            tx,
+            resource: make_resource("test-service".to_owned()),
+        };
+        let metrics = vec![Metric::default(); MAX_METRICS_PER_DONATION + 1];
+        exporter.export(metrics);
+
+        let mut payloads = 0;
+        let mut total = 0;
+        while let Ok(PumpMsg::Payload(p)) = rx.try_recv() {
+            payloads += 1;
+            total += decode(p).resource_metrics[0].scope_metrics[0].metrics.len();
+        }
+        assert_eq!(payloads, 2, "one full chunk + remainder");
+        assert_eq!(total, MAX_METRICS_PER_DONATION + 1);
+    }
+
+    #[test]
+    fn summary_and_untyped_families_are_skipped() {
+        // An empty registry gathers nothing; convert yields nothing.
+        let registry = Registry::new();
+        assert!(convert_families(&registry.gather()).is_empty());
+    }
+}

crates/common/xai-computer-hub-sdk/src/metrics.rs

@@ -0,0 +1,552 @@
+//! Feature-gated Prometheus metrics for the SDK.
+//!
+//! When the `metrics` cargo feature is enabled, each helper records to a
+//! lazily-registered Prometheus counter / gauge / histogram. When
+//! disabled (the default), every helper compiles to an empty function
+//! body so the SDK carries zero prometheus dependency.
+
+#[cfg(feature = "metrics")]
+mod inner {
+    use prometheus::{
+        Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec,
+        exponential_buckets, register_histogram, register_histogram_vec, register_int_counter,
+        register_int_counter_vec, register_int_gauge, register_int_gauge_vec,
+    };
+    use std::sync::LazyLock;
+
+    static POOL_CONNECTIONS: LazyLock<IntGauge> = LazyLock::new(|| {
+        register_int_gauge!(
+            "computer_hub_client_pool_connections",
+            "Active pooled connections in the SDK connection pool."
+        )
+        .expect("computer_hub_client_pool_connections must register once")
+    });
+
+    static POOL_EVICTIONS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "computer_hub_client_pool_evictions_total",
+            "Pooled connections closed by the idle reaper (unused past the idle TTL)."
+        )
+        .expect("computer_hub_client_pool_evictions_total must register once")
+    });
+
+    static RECONNECTS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "computer_hub_client_reconnects_total",
+            "Cumulative reconnect attempts that succeeded."
+        )
+        .expect("computer_hub_client_reconnects_total must register once")
+    });
+
+    static RECONNECT_FAILED_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
+        register_int_counter_vec!(
+            "computer_hub_client_reconnect_failed_total",
+            "Cumulative reconnect attempts that failed, by reason \
+             (handshake_auth = fatal 401/403, transport = retryable).",
+            &["reason"]
+        )
+        .expect("computer_hub_client_reconnect_failed_total must register once")
+    });
+
+    static RECONNECT_DURATION_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+        register_histogram!(
+            "computer_hub_client_reconnect_duration_seconds",
+            "Time to complete a reconnect cycle (handshake + session/tool replay).",
+            exponential_buckets(0.01, 2.0, 14).expect("valid bucket params")
+        )
+        .expect("computer_hub_client_reconnect_duration_seconds must register once")
+    });
+
+    static RECONNECTS_BY_CAUSE_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
+        register_int_counter_vec!(
+            "computer_hub_client_reconnects_by_cause_total",
+            "Successful reconnects by disconnect cause of the previous connection \
+             (close_frame, eof, transport_read_error, transport_write_error, forced). \
+             Cause-labeled companion to computer_hub_client_reconnects_total.",
+            &["cause"]
+        )
+        .expect("computer_hub_client_reconnects_by_cause_total must register once")
+    });
+
+    static RECONNECT_GAP_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+        register_histogram!(
+            "computer_hub_client_reconnect_gap_seconds",
+            "Time from the last inbound frame on the dead connection to a successful reconnect.",
+            exponential_buckets(0.1, 2.0, 14).expect("valid bucket params")
+        )
+        .expect("computer_hub_client_reconnect_gap_seconds must register once")
+    });
+
+    static CALL_DISPATCH_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+        register_histogram!(
+            "computer_hub_client_call_dispatch_seconds",
+            "Time to set up and queue the outbound remote dispatch.",
+            exponential_buckets(0.0001, 2.0, 14).expect("valid bucket params")
+        )
+        .expect("computer_hub_client_call_dispatch_seconds must register once")
+    });
+
+    static DEMUX_INBOX_DEPTH: LazyLock<IntGauge> = LazyLock::new(|| {
+        register_int_gauge!(
+            "computer_hub_client_demux_inbox_depth",
+            "Number of session inboxes registered in the inbound demux."
+        )
+        .expect("computer_hub_client_demux_inbox_depth must register once")
+    });
+
+    static CALL_ID_COLLISIONS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "computer_hub_client_call_id_collisions_total",
+            "Call-id collisions detected in the harness dispatch path."
+        )
+        .expect("computer_hub_client_call_id_collisions_total must register once")
+    });
+
+    // ── Server integration metrics ──────────────────────────────────
+
+    static HARNESS_CONNECT_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
+        register_int_counter_vec!(
+            "hub_harness_connect_total",
+            "Hub connection attempts by outcome and sampler.",
+            &["status", "sampler"]
+        )
+        .expect("hub_harness_connect_total must register once")
+    });
+
+    static SESSION_EVENT_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
+        register_int_counter_vec!(
+            "hub_session_event_total",
+            "SessionEvent emissions by event type.",
+            &["event_type"]
+        )
+        .expect("hub_session_event_total must register once")
+    });
+
+    static SESSION_OP_DURATION_SECONDS: LazyLock<HistogramVec> = LazyLock::new(|| {
+        register_histogram_vec!(
+            "hub_session_op_duration_seconds",
+            "Latency of hub session lifecycle operations (open/bind) by op and outcome.",
+            &["op", "status"],
+            exponential_buckets(0.001, 2.0, 14).expect("valid bucket params")
+        )
+        .expect("hub_session_op_duration_seconds must register once")
+    });
+
+    static SESSION_SOFT_REBIND_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_session_soft_rebind_total",
+            "Redundant session.bind frames for a session with a live dispatch loop, \
+             handled as a non-destructive soft rebind (serve state refreshed, \
+             in-flight calls preserved)."
+        )
+        .expect("hub_session_soft_rebind_total must register once")
+    });
+
+    static NO_HANDLER_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_sdk_no_handler_total",
+            "tool_call_request frames rejected with -32011 because the session's \
+             current handler set has no handler for the requested tool_id."
+        )
+        .expect("hub_sdk_no_handler_total must register once")
+    });
+
+    static HOOK_SEND_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
+        register_int_counter_vec!(
+            "hub_hook_send_total",
+            "Hook sends by hook type.",
+            &["hook_type"]
+        )
+        .expect("hub_hook_send_total must register once")
+    });
+
+    static PROGRESS_FRAMES_FORWARDED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_progress_frames_forwarded_total",
+            "Progress frames forwarded by ToolServer."
+        )
+        .expect("hub_progress_frames_forwarded_total must register once")
+    });
+
+    static CANCEL_HOOK_RECEIVED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_cancel_hook_received_total",
+            "Cancel hooks received by workspace tool server."
+        )
+        .expect("hub_cancel_hook_received_total must register once")
+    });
+
+    static WRITER_SINK_SEND_ERRORS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "computer_hub_client_writer_sink_send_errors_total",
+            "Writer-task sink send failures; each signals the reader to reconnect."
+        )
+        .expect("computer_hub_client_writer_sink_send_errors_total must register once")
+    });
+
+    static RECONNECT_WRITER_RESUME_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "computer_hub_client_reconnect_writer_resume_total",
+            "Fresh-sink Resume handoffs delivered to the writer task after reconnect."
+        )
+        .expect("computer_hub_client_reconnect_writer_resume_total must register once")
+    });
+
+    static LIVENESS_DEADLINE_EXPIRED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "computer_hub_client_liveness_deadline_expired_total",
+            "Liveness-deadline expiries in the reader (no inbound WebSocket \
+             frame within the deadline); each declares the connection dead and \
+             drives the normal reconnect path."
+        )
+        .expect("computer_hub_client_liveness_deadline_expired_total must register once")
+    });
+
+    static HEARTBEAT_PONG_DROPPED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "computer_hub_client_heartbeat_pong_dropped_total",
+            "App-level heartbeat pongs dropped because outbound_tx was saturated (split reader)."
+        )
+        .expect("computer_hub_client_heartbeat_pong_dropped_total must register once")
+    });
+
+    static CANCEL_APPLIED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_cancel_applied_total",
+            "Cancel hooks that hit a live in-flight call and cancelled it."
+        )
+        .expect("hub_cancel_applied_total must register once")
+    });
+
+    static CANCEL_PENDING_TOMBSTONED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_cancel_pending_tombstoned_total",
+            "Cancel hooks recorded as a pending tombstone (call not yet registered or already done)."
+        )
+        .expect("hub_cancel_pending_tombstoned_total must register once")
+    });
+
+    static CANCEL_NO_TARGET_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_cancel_no_target_total",
+            "Cancel hooks with no call_id (session-wide, no specific call to cancel)."
+        )
+        .expect("hub_cancel_no_target_total must register once")
+    });
+
+    static TOOL_CALL_REJECTED_OVERLOADED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_tool_call_rejected_overloaded_total",
+            "Tool calls rejected by admission timeout (-32016 tool_busy)."
+        )
+        .expect("hub_tool_call_rejected_overloaded_total must register once")
+    });
+
+    static INBOX_FULL_REQUEST_REJECTED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_inbox_full_request_rejected_total",
+            "Requests rejected with an overloaded response on a full session inbox."
+        )
+        .expect("hub_inbox_full_request_rejected_total must register once")
+    });
+
+    static INBOX_FULL_REJECT_SEND_FAILED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_inbox_full_reject_send_failed_total",
+            "Overloaded rejections dropped because outbound was also full (residual silent loss)."
+        )
+        .expect("hub_inbox_full_reject_send_failed_total must register once")
+    });
+
+    static INBOX_FULL_NOTIFICATION_DROPPED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_inbox_full_notification_dropped_total",
+            "Notifications (no id) dropped on a full session inbox."
+        )
+        .expect("hub_inbox_full_notification_dropped_total must register once")
+    });
+
+    static SERVE_REPLAY_TIMEOUT_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "computer_hub_client_serve_replay_timeout_total",
+            "serve attempts that hit the per-attempt reply deadline, from any \
+             serve call site (reconnect replay, run(), bind, tool updates)."
+        )
+        .expect("computer_hub_client_serve_replay_timeout_total must register once")
+    });
+
+    static NOTIF_LAGGED_RECOVERED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_notif_lagged_recovered_total",
+            "Connection-notification broadcast Lagged events recovered by \
+             continuing the loop instead of exiting."
+        )
+        .expect("hub_notif_lagged_recovered_total must register once")
+    });
+
+    static EARLY_NOTIF_BUFFERED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
+        register_int_counter!(
+            "hub_early_notif_buffered_total",
+            "Connection-level notification frames (binds, unbinds, evicts, \
+             ...) buffered between connect and ToolServer::run() and replayed \
+             instead of dropped."
+        )
+        .expect("hub_early_notif_buffered_total must register once")
+    });
+
+    static TOOL_CALL_INFLIGHT: LazyLock<IntGaugeVec> = LazyLock::new(|| {
+        register_int_gauge_vec!(
+            "hub_tool_call_inflight",
+            "Concurrent running tool calls holding an admission permit, by scope.",
+            &["scope"]
+        )
+        .expect("hub_tool_call_inflight must register once")
+    });
+
+    static ADMISSION_WAIT_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
+        register_histogram!(
+            "hub_tool_call_admission_wait_seconds",
+            "Time blocked acquiring the three admission permits (one shared deadline).",
+            exponential_buckets(0.0001, 2.0, 16).expect("valid bucket params")
+        )
+        .expect("hub_tool_call_admission_wait_seconds must register once")
+    });
+
+    pub(crate) fn pool_connections_inc() {
+        POOL_CONNECTIONS.inc();
+    }
+
+    pub(crate) fn pool_connections_dec() {
+        POOL_CONNECTIONS.dec();
+    }
+
+    pub(crate) fn pool_evictions_inc() {
+        POOL_EVICTIONS_TOTAL.inc();
+    }
+
+    pub(crate) fn reconnect_succeeded() {
+        RECONNECTS_TOTAL.inc();
+    }
+
+    pub(crate) fn reconnect_failed(reason: &str) {
+        RECONNECT_FAILED_TOTAL.with_label_values(&[reason]).inc();
+    }
+
+    pub(crate) fn reconnect_duration_observe(secs: f64) {
+        RECONNECT_DURATION_SECONDS.observe(secs);
+    }
+
+    pub(crate) fn reconnect_cause(cause: &str) {
+        RECONNECTS_BY_CAUSE_TOTAL.with_label_values(&[cause]).inc();
+    }
+
+    pub(crate) fn reconnect_gap_observe(secs: f64) {
+        RECONNECT_GAP_SECONDS.observe(secs);
+    }
+
+    pub(crate) fn call_dispatch_observe(secs: f64) {
+        CALL_DISPATCH_SECONDS.observe(secs);
+    }
+
+    pub(crate) fn demux_inbox_depth_set(depth: i64) {
+        DEMUX_INBOX_DEPTH.set(depth);
+    }
+
+    pub(crate) fn call_id_collision() {
+        CALL_ID_COLLISIONS_TOTAL.inc();
+    }
+
+    /// Record a harness connection attempt.
+    ///
+    /// `sampler` identifies the caller (`"chat"` or `"shell"`).
+    /// `status` is `"ok"`, `"error"`, or `"fallback"` (fallback is
+    /// emitted by the caller in `AgentBuilder::build_harness()`, not
+    /// by the SDK).
+    pub fn harness_connect(status: &str, sampler: &str) {
+        HARNESS_CONNECT_TOTAL
+            .with_label_values(&[status, sampler])
+            .inc();
+    }
+
+    pub(crate) fn session_event(event_type: &str) {
+        SESSION_EVENT_TOTAL.with_label_values(&[event_type]).inc();
+    }
+
+    /// Observe the latency of a session lifecycle operation.
+    /// `op` is `"open"` or `"bind"`; `status` is `"ok"` or `"error"`.
+    pub(crate) fn session_op_observe(op: &str, status: &str, secs: f64) {
+        SESSION_OP_DURATION_SECONDS
+            .with_label_values(&[op, status])
+            .observe(secs);
+    }
+
+    pub(crate) fn session_soft_rebind() {
+        SESSION_SOFT_REBIND_TOTAL.inc();
+    }
+
+    pub(crate) fn no_handler() {
+        NO_HANDLER_TOTAL.inc();
+    }
+
+    pub(crate) fn hook_send(hook_type: &str) {
+        HOOK_SEND_TOTAL.with_label_values(&[hook_type]).inc();
+    }
+
+    pub(crate) fn progress_frame_forwarded() {
+        PROGRESS_FRAMES_FORWARDED_TOTAL.inc();
+    }
+
+    pub(crate) fn cancel_hook_received() {
+        CANCEL_HOOK_RECEIVED_TOTAL.inc();
+    }
+
+    pub(crate) fn writer_sink_send_error() {
+        WRITER_SINK_SEND_ERRORS_TOTAL.inc();
+    }
+
+    pub(crate) fn reconnect_writer_resume() {
+        RECONNECT_WRITER_RESUME_TOTAL.inc();
+    }
+
+    pub(crate) fn liveness_deadline_expired() {
+        LIVENESS_DEADLINE_EXPIRED_TOTAL.inc();
+    }
+
+    pub(crate) fn heartbeat_pong_dropped() {
+        HEARTBEAT_PONG_DROPPED_TOTAL.inc();
+    }
+
+    pub(crate) fn cancel_applied() {
+        CANCEL_APPLIED_TOTAL.inc();
+    }
+
+    pub(crate) fn cancel_pending_tombstoned() {
+        CANCEL_PENDING_TOMBSTONED_TOTAL.inc();
+    }
+
+    pub(crate) fn cancel_no_target() {
+        CANCEL_NO_TARGET_TOTAL.inc();
+    }
+
+    pub(crate) fn tool_call_rejected_overloaded() {
+        TOOL_CALL_REJECTED_OVERLOADED_TOTAL.inc();
+    }
+
+    pub(crate) fn inbox_full_request_rejected() {
+        INBOX_FULL_REQUEST_REJECTED_TOTAL.inc();
+    }
+
+    pub(crate) fn inbox_full_reject_send_failed() {
+        INBOX_FULL_REJECT_SEND_FAILED_TOTAL.inc();
+    }
+
+    pub(crate) fn inbox_full_notification_dropped() {
+        INBOX_FULL_NOTIFICATION_DROPPED_TOTAL.inc();
+    }
+
+    pub(crate) fn serve_replay_timeout() {
+        SERVE_REPLAY_TIMEOUT_TOTAL.inc();
+    }
+
+    pub(crate) fn notif_lagged_recovered() {
+        NOTIF_LAGGED_RECOVERED_TOTAL.inc();
+    }
+
+    pub(crate) fn early_notif_buffered(frames: u64) {
+        EARLY_NOTIF_BUFFERED_TOTAL.inc_by(frames);
+    }
+
+    pub(crate) fn tool_call_inflight_inc(scope: &str) {
+        TOOL_CALL_INFLIGHT.with_label_values(&[scope]).inc();
+    }
+
+    pub(crate) fn tool_call_inflight_dec(scope: &str) {
+        TOOL_CALL_INFLIGHT.with_label_values(&[scope]).dec();
+    }
+
+    pub(crate) fn admission_wait_observe(secs: f64) {
+        ADMISSION_WAIT_SECONDS.observe(secs);
+    }
+}
+
+#[cfg(not(feature = "metrics"))]
+mod inner {
+    pub(crate) fn pool_connections_inc() {}
+    pub(crate) fn pool_connections_dec() {}
+    pub(crate) fn pool_evictions_inc() {}
+    pub(crate) fn reconnect_succeeded() {}
+    pub(crate) fn reconnect_failed(_reason: &str) {}
+    pub(crate) fn reconnect_duration_observe(_secs: f64) {}
+    pub(crate) fn reconnect_cause(_cause: &str) {}
+    pub(crate) fn reconnect_gap_observe(_secs: f64) {}
+    pub(crate) fn call_dispatch_observe(_secs: f64) {}
+    pub(crate) fn demux_inbox_depth_set(_depth: i64) {}
+    pub(crate) fn call_id_collision() {}
+    pub fn harness_connect(_status: &str, _sampler: &str) {}
+    pub(crate) fn session_event(_event_type: &str) {}
+    pub(crate) fn session_op_observe(_op: &str, _status: &str, _secs: f64) {}
+    pub(crate) fn session_soft_rebind() {}
+    pub(crate) fn no_handler() {}
+    pub(crate) fn hook_send(_hook_type: &str) {}
+    pub(crate) fn progress_frame_forwarded() {}
+    pub(crate) fn cancel_hook_received() {}
+    pub(crate) fn writer_sink_send_error() {}
+    pub(crate) fn reconnect_writer_resume() {}
+    pub(crate) fn liveness_deadline_expired() {}
+    pub(crate) fn heartbeat_pong_dropped() {}
+    pub(crate) fn cancel_applied() {}
+    pub(crate) fn cancel_pending_tombstoned() {}
+    pub(crate) fn cancel_no_target() {}
+    pub(crate) fn tool_call_rejected_overloaded() {}
+    pub(crate) fn inbox_full_request_rejected() {}
+    pub(crate) fn inbox_full_reject_send_failed() {}
+    pub(crate) fn inbox_full_notification_dropped() {}
+    pub(crate) fn serve_replay_timeout() {}
+    pub(crate) fn notif_lagged_recovered() {}
+    pub(crate) fn early_notif_buffered(_frames: u64) {}
+    pub(crate) fn tool_call_inflight_inc(_scope: &str) {}
+    pub(crate) fn tool_call_inflight_dec(_scope: &str) {}
+    pub(crate) fn admission_wait_observe(_secs: f64) {}
+}
+
+pub(crate) use inner::admission_wait_observe;
+pub(crate) use inner::call_dispatch_observe;
+pub(crate) use inner::call_id_collision;
+pub(crate) use inner::cancel_applied;
+pub(crate) use inner::cancel_hook_received;
+pub(crate) use inner::cancel_no_target;
+pub(crate) use inner::cancel_pending_tombstoned;
+pub(crate) use inner::demux_inbox_depth_set;
+pub(crate) use inner::early_notif_buffered;
+pub(crate) use inner::heartbeat_pong_dropped;
+pub(crate) use inner::hook_send;
+pub(crate) use inner::inbox_full_notification_dropped;
+pub(crate) use inner::inbox_full_reject_send_failed;
+pub(crate) use inner::inbox_full_request_rejected;
+pub(crate) use inner::liveness_deadline_expired;
+pub(crate) use inner::no_handler;
+pub(crate) use inner::notif_lagged_recovered;
+pub(crate) use inner::pool_connections_dec;
+pub(crate) use inner::pool_connections_inc;
+pub(crate) use inner::pool_evictions_inc;
+pub(crate) use inner::progress_frame_forwarded;
+pub(crate) use inner::reconnect_cause;
+pub(crate) use inner::reconnect_duration_observe;
+pub(crate) use inner::reconnect_failed;
+pub(crate) use inner::reconnect_gap_observe;
+pub(crate) use inner::reconnect_succeeded;
+pub(crate) use inner::reconnect_writer_resume;
+pub(crate) use inner::serve_replay_timeout;
+pub(crate) use inner::session_event;
+pub(crate) use inner::session_op_observe;
+pub(crate) use inner::session_soft_rebind;
+pub(crate) use inner::tool_call_inflight_dec;
+pub(crate) use inner::tool_call_inflight_inc;
+pub(crate) use inner::tool_call_rejected_overloaded;
+pub(crate) use inner::writer_sink_send_error;
+
+/// Record a harness connection attempt. Public so callers outside
+/// the SDK (e.g. `AgentBuilder::build_harness()` in the agentic sampler)
+/// can emit `status="fallback"` when the server connection fails and the
+/// builder falls back to a local-only harness.
+pub use inner::harness_connect;

crates/common/xai-computer-hub-sdk/src/notification.rs

@@ -0,0 +1,362 @@
+//! Parsed server notification events.
+//!
+//! [`HubNotification`] is the typed representation of server-pushed
+//! notification frames that arrive on a session inbox. The
+//! [`HubNotification::parse`] constructor classifies a raw JSON value
+//! by its `method` field and deserializes the known shapes; anything
+//! unrecognised lands in [`HubNotification::Unknown`] so callers never
+//! lose data.
+
+use serde_json::Value;
+use tracing::warn;
+use xai_tool_protocol::{
+    SessionId, ToolId, ToolNotificationFrame, ToolServerStatusPayload, ToolsChanged,
+};
+
+/// A typed server notification event parsed from a raw JSON-RPC notification frame.
+#[derive(Debug, Clone, PartialEq)]
+pub enum HubNotification {
+    /// The active tool set for a session changed (tools added, removed, or updated).
+    ToolsChanged {
+        session_id: SessionId,
+        added: Vec<ToolId>,
+        removed: Vec<ToolId>,
+        updated: Vec<ToolId>,
+    },
+    /// A tool notification forwarded by the server to all subscribers.
+    ToolNotification {
+        session_id: SessionId,
+        frame: ToolNotificationFrame,
+    },
+    /// Tool server lifecycle status change, extracted from
+    /// `__tool_server_status` / `status_changed` notification frames.
+    ToolServerStatusChanged {
+        session_id: SessionId,
+        status: ToolServerStatusPayload,
+    },
+    /// A notification whose `method` is not recognised by this SDK version.
+    Unknown { method: String, params: Value },
+}
+
+impl HubNotification {
+    /// Parse a raw JSON-RPC notification into a typed [`HubNotification`].
+    ///
+    /// Returns `None` when the value lacks a `method` field (i.e. it is
+    /// not a notification at all).
+    pub fn parse(value: &Value) -> Option<Self> {
+        let method = value.get("method")?.as_str()?;
+        let params = value
+            .get("params")
+            .cloned()
+            .unwrap_or(Value::Object(Default::default()));
+
+        match method {
+            // `ToolsChanged` carries `session_id` inside `params`.
+            "tools_changed" => match serde_json::from_value::<ToolsChanged>(params.clone()) {
+                Ok(tc) => Some(HubNotification::ToolsChanged {
+                    session_id: tc.session_id,
+                    added: tc.added,
+                    removed: tc.removed,
+                    updated: tc.updated,
+                }),
+                Err(err) => {
+                    warn!(%err, "tools_changed params failed to deserialize; falling back to Unknown");
+                    Some(HubNotification::Unknown {
+                        method: method.to_owned(),
+                        params,
+                    })
+                }
+            },
+            // `ToolNotificationFrame` has no `session_id`; use the envelope field.
+            "tool.notification" => {
+                let frame_result = serde_json::from_value::<ToolNotificationFrame>(params.clone());
+                let session_id = value
+                    .get("session_id")
+                    .and_then(Value::as_str)
+                    .and_then(|s| SessionId::new(s).ok());
+                match (frame_result, session_id) {
+                    (Ok(frame), Some(session_id)) => {
+                        if frame
+                            .tool_id
+                            .as_ref()
+                            .is_some_and(|id| id.as_str() == "__tool_server_status")
+                            && let xai_tool_protocol::notification_wire::WireToolNotification::Custom(ref c) = frame.notification
+                            && c.kind == "status_changed"
+                        {
+                            match serde_json::from_value::<ToolServerStatusPayload>(
+                                c.payload.clone(),
+                            ) {
+                                Ok(status) => {
+                                    return Some(HubNotification::ToolServerStatusChanged {
+                                        session_id,
+                                        status,
+                                    });
+                                }
+                                Err(err) => {
+                                    warn!(%err, "tool_server status payload failed to deserialize");
+                                }
+                            }
+                        }
+                        Some(HubNotification::ToolNotification { session_id, frame })
+                    }
+                    (Err(err), _) => {
+                        warn!(%err, "tool.notification params failed to deserialize; falling back to Unknown");
+                        Some(HubNotification::Unknown {
+                            method: method.to_owned(),
+                            params,
+                        })
+                    }
+                    (_, None) => {
+                        warn!(
+                            "tool.notification missing or invalid session_id; falling back to Unknown"
+                        );
+                        Some(HubNotification::Unknown {
+                            method: method.to_owned(),
+                            params,
+                        })
+                    }
+                }
+            }
+            _ => Some(HubNotification::Unknown {
+                method: method.to_owned(),
+                params,
+            }),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn parse_tools_changed() {
+        let value = json!({
+            "jsonrpc": "2.0",
+            "session_id": "s1",
+            "method": "tools_changed",
+            "params": {
+                "session_id": "s1",
+                "added": ["echo", "add"],
+                "removed": [],
+            }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse");
+        match notif {
+            HubNotification::ToolsChanged {
+                session_id,
+                added,
+                removed,
+                updated,
+            } => {
+                assert_eq!(session_id.as_str(), "s1");
+                assert_eq!(added.len(), 2);
+                assert!(removed.is_empty());
+                assert!(updated.is_empty());
+            }
+            other => panic!("expected ToolsChanged, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn parse_tools_changed_with_updated() {
+        let value = json!({
+            "jsonrpc": "2.0",
+            "session_id": "s1",
+            "method": "tools_changed",
+            "params": {
+                "session_id": "s1",
+                "added": ["new_tool"],
+                "removed": ["old_tool"],
+                "updated": ["echo", "add"],
+            }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse");
+        match notif {
+            HubNotification::ToolsChanged {
+                session_id,
+                added,
+                removed,
+                updated,
+            } => {
+                assert_eq!(session_id.as_str(), "s1");
+                assert_eq!(added.len(), 1);
+                assert_eq!(removed.len(), 1);
+                assert_eq!(updated.len(), 2);
+                assert_eq!(updated[0].as_str(), "echo");
+                assert_eq!(updated[1].as_str(), "add");
+            }
+            other => panic!("expected ToolsChanged, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn parse_tool_notification_custom() {
+        let value = json!({
+            "jsonrpc": "2.0",
+            "session_id": "s1",
+            "method": "tool.notification",
+            "params": {
+                "tool_id": "echo",
+                "notification": {
+                    "shape": "custom",
+                    "value": {
+                        "kind": "echo.status",
+                        "payload": { "status": "idle" }
+                    }
+                }
+            }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse");
+        match notif {
+            HubNotification::ToolNotification { session_id, frame } => {
+                assert_eq!(session_id.as_str(), "s1");
+                assert_eq!(frame.tool_id.as_ref().unwrap().as_str(), "echo");
+            }
+            other => panic!("expected ToolNotification, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn parse_tool_notification_missing_session_id_falls_back_to_unknown() {
+        let value = json!({
+            "jsonrpc": "2.0",
+            "method": "tool.notification",
+            "params": {
+                "tool_id": "echo",
+                "notification": {
+                    "shape": "custom",
+                    "value": { "kind": "test", "payload": {} }
+                }
+            }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None");
+        assert!(
+            matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tool.notification"),
+            "tool.notification without envelope session_id should fall back to Unknown, got {notif:?}"
+        );
+    }
+
+    #[test]
+    fn parse_unknown_method() {
+        let value = json!({
+            "jsonrpc": "2.0",
+            "session_id": "s1",
+            "method": "future.method",
+            "params": { "key": "value" }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse");
+        match notif {
+            HubNotification::Unknown { method, params } => {
+                assert_eq!(method, "future.method");
+                assert_eq!(params["key"], "value");
+            }
+            other => panic!("expected Unknown, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn parse_missing_method_returns_none() {
+        let value = json!({ "jsonrpc": "2.0", "id": "123", "result": {} });
+        assert!(HubNotification::parse(&value).is_none());
+    }
+
+    #[test]
+    fn parse_tools_changed_bad_params_falls_back_to_unknown() {
+        // `params` has wrong shape (missing required fields) — should fall
+        // back to Unknown instead of returning None and dropping the event.
+        let value = json!({
+            "jsonrpc": "2.0",
+            "method": "tools_changed",
+            "params": { "unexpected_field": true }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None");
+        assert!(
+            matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tools_changed"),
+            "malformed tools_changed should fall back to Unknown, got {notif:?}"
+        );
+    }
+
+    #[test]
+    fn parse_tool_server_status_changed() {
+        let value = json!({
+            "jsonrpc": "2.0",
+            "session_id": "s1",
+            "method": "tool.notification",
+            "params": {
+                "tool_id": "__tool_server_status",
+                "notification": {
+                    "shape": "custom",
+                    "value": {
+                        "kind": "status_changed",
+                        "payload": {
+                            "status": "busy",
+                            "active_tool_calls": 2,
+                            "active_tool_names": ["read_file", "grep"],
+                            "background_tasks": 0,
+                            "pending_tool_calls": 0,
+                            "last_tool_call_started_ms": 100,
+                            "last_tool_call_completed_ms": 0,
+                            "uptime_ms": 5000,
+                        }
+                    }
+                }
+            }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse");
+        match notif {
+            HubNotification::ToolServerStatusChanged { session_id, status } => {
+                assert_eq!(session_id.as_str(), "s1");
+                assert_eq!(
+                    status.status,
+                    xai_tool_protocol::ToolServerLifecycleStatus::Busy
+                );
+                assert_eq!(status.active_tool_calls, 2);
+            }
+            other => panic!("expected ToolServerStatusChanged, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn parse_tool_server_status_non_status_tool_id_stays_generic() {
+        // A tool.notification with a different tool_id should remain
+        // as ToolNotification, not be intercepted.
+        let value = json!({
+            "jsonrpc": "2.0",
+            "session_id": "s1",
+            "method": "tool.notification",
+            "params": {
+                "tool_id": "some_other_tool",
+                "notification": {
+                    "shape": "custom",
+                    "value": {
+                        "kind": "status_changed",
+                        "payload": { "status": "ready" }
+                    }
+                }
+            }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse");
+        assert!(
+            matches!(notif, HubNotification::ToolNotification { .. }),
+            "non-__tool_server_status tool_id should stay as ToolNotification, got {notif:?}"
+        );
+    }
+
+    #[test]
+    fn parse_tool_notification_bad_params_falls_back_to_unknown() {
+        // `params` has wrong shape — should fall back to Unknown.
+        let value = json!({
+            "jsonrpc": "2.0",
+            "session_id": "s1",
+            "method": "tool.notification",
+            "params": { "not_a_valid_frame": true }
+        });
+        let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None");
+        assert!(
+            matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tool.notification"),
+            "malformed tool.notification should fall back to Unknown, got {notif:?}"
+        );
+    }
+}

crates/common/xai-computer-hub-sdk/src/observability.rs

@@ -0,0 +1,275 @@
+//! Server-side session event emitter.
+//!
+//! [`ObservabilityBridge`] is a thin facade for emitting session-level
+//! events (turn lifecycle, phase changes) to the connected server. Tool-call events
+//! (`ToolCallStarted` / `ToolCallCompleted`) are emitted automatically
+//! by [`crate::harness::ToolHarness::call`] and do not need the bridge.
+//!
+//! The caller is responsible for also emitting to the local sink
+//! (`EventTracker` in the shell, `EventProcPublisher` in the
+//! chat service) — the bridge handles only the server leg.
+//!
+//! This separation is deliberate: each sampler's local sink has a
+//! different type and API surface.  Forcing a trait/callback into the
+//! bridge would add abstraction overhead without benefit, since the
+//! call sites already have the local sink in scope.
+
+use std::sync::Arc;
+
+use xai_tool_protocol::{SessionId, session_event::SessionEvent};
+
+use crate::harness::ToolHarness;
+
+/// Emits [`SessionEvent`]s to the connected server as `ToolNotificationFrame` custom
+/// notifications with `kind = "session_event"`.
+///
+/// No-ops gracefully when no harness is present (i.e. `harness` is
+/// `None`).  Server notification failures are silently ignored — the bridge
+/// is fire-and-forget so server issues never affect the sampler's main loop.
+///
+/// Callers MUST also emit to their local sink separately:
+/// - Shell: `self.events.emit(Event::...)`
+/// - Chat service: `publisher.publish_agent_event(...)`
+pub struct ObservabilityBridge {
+    harness: Option<Arc<ToolHarness>>,
+    /// Retained for future payload enrichment and logging.
+    session_id: SessionId,
+}
+
+impl ObservabilityBridge {
+    pub fn new(harness: Option<Arc<ToolHarness>>, session_id: SessionId) -> Self {
+        Self {
+            harness,
+            session_id,
+        }
+    }
+
+    /// The session id this bridge was created for.
+    pub fn session_id(&self) -> &SessionId {
+        &self.session_id
+    }
+
+    /// Whether a harness is present (i.e. server emission is active).
+    pub fn has_harness(&self) -> bool {
+        self.harness.is_some()
+    }
+
+    /// Emit a session event to the connected server.  No-ops if no harness is present.
+    ///
+    /// Delegates frame construction + wire dispatch to
+    /// [`ToolHarness::emit_session_event`] so the SDK keeps a single
+    /// canonical encoding path.
+    ///
+    /// Callers MUST also emit to their local sink separately:
+    /// - Shell: `self.events.emit(Event::...)`
+    /// - Chat service: `publisher.publish_agent_event(...)`
+    pub async fn emit(&self, event: SessionEvent) {
+        let event_type = match &event {
+            SessionEvent::TurnStarted { .. } => "turn_started",
+            SessionEvent::TurnEnded { .. } => "turn_ended",
+            SessionEvent::ToolCallStarted { .. } => "tool_call_started",
+            SessionEvent::ToolCallCompleted { .. } => "tool_call_completed",
+            SessionEvent::PhaseChanged { .. } => "phase_changed",
+            SessionEvent::Unknown => "unknown",
+        };
+        crate::metrics::session_event(event_type);
+        if let Some(harness) = &self.harness {
+            harness.emit_session_event(event).await;
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use xai_tool_protocol::session_event::{SessionEvent, SessionPhase, ToolCallOutcome};
+    use xai_tool_protocol::turn_hook::TurnHookOutcome;
+
+    fn test_session_id() -> SessionId {
+        SessionId::new("test-obs-session").expect("valid")
+    }
+
+    // ── No-harness path ─────────────────────────────────────────────
+
+    #[tokio::test]
+    async fn emit_without_harness_is_noop() {
+        let bridge = ObservabilityBridge::new(None, test_session_id());
+        // Must not panic and should return immediately.
+        bridge
+            .emit(SessionEvent::TurnStarted {
+                turn_number: 1,
+                model_id: "grok-3".into(),
+                yolo_mode: false,
+            })
+            .await;
+    }
+
+    #[test]
+    fn has_harness_returns_false_when_none() {
+        let bridge = ObservabilityBridge::new(None, test_session_id());
+        assert!(!bridge.has_harness());
+    }
+
+    // ── Constructor field storage ───────────────────────────────────
+
+    #[test]
+    fn new_stores_session_id() {
+        let sid = test_session_id();
+        let bridge = ObservabilityBridge::new(None, sid.clone());
+        assert_eq!(bridge.session_id(), &sid);
+    }
+
+    #[test]
+    fn has_harness_returns_true_when_present() {
+        let harness = ToolHarness::local_only_with(
+            crate::harness::LocalRegistry::new(),
+            test_session_id(),
+            Default::default(),
+        );
+        let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id());
+        assert!(bridge.has_harness());
+    }
+
+    // ── Serialization correctness ───────────────────────────────────
+
+    #[test]
+    fn session_event_serializes_to_expected_json() {
+        let event = SessionEvent::TurnStarted {
+            turn_number: 1,
+            model_id: "grok-3".into(),
+            yolo_mode: true,
+        };
+        let value = serde_json::to_value(&event).unwrap();
+        assert_eq!(value["event_type"], "turn_started");
+        assert_eq!(value["turn_number"], 1);
+        assert_eq!(value["model_id"], "grok-3");
+        assert_eq!(value["yolo_mode"], true);
+    }
+
+    #[test]
+    fn session_event_turn_ended_serializes_correctly() {
+        let event = SessionEvent::TurnEnded {
+            turn_number: 5,
+            outcome: TurnHookOutcome::Completed,
+            duration_ms: 3200,
+            tool_call_count: 12,
+            model_id: "grok-3".into(),
+        };
+        let value = serde_json::to_value(&event).unwrap();
+        assert_eq!(value["event_type"], "turn_ended");
+        assert_eq!(value["outcome"], "completed");
+        assert_eq!(value["tool_call_count"], 12);
+    }
+
+    #[test]
+    fn session_event_tool_call_completed_serializes_correctly() {
+        let event = SessionEvent::ToolCallCompleted {
+            tool_call_id: "call-1".into(),
+            tool_name: "bash".into(),
+            duration_ms: 500,
+            outcome: ToolCallOutcome::Success,
+        };
+        let value = serde_json::to_value(&event).unwrap();
+        assert_eq!(value["event_type"], "tool_call_completed");
+        assert_eq!(value["outcome"], "success");
+    }
+
+    #[test]
+    fn session_event_phase_changed_serializes_correctly() {
+        let event = SessionEvent::PhaseChanged {
+            phase: SessionPhase::Sampling,
+        };
+        let value = serde_json::to_value(&event).unwrap();
+        assert_eq!(value["event_type"], "phase_changed");
+        assert_eq!(value["phase"], "sampling");
+    }
+
+    // Frame-construction invariants moved to `crate::harness` where the
+    // builder now lives (`ToolHarness::emit_session_event`).
+
+    // ── End-to-end with local-only harness ───────────────────────────
+
+    #[tokio::test]
+    async fn emit_with_local_only_harness_does_not_panic() {
+        // A local-only harness has no server connection, so
+        // `send_notification` returns `Err` — but the bridge ignores
+        // errors, so this must succeed silently.
+        let harness = ToolHarness::local_only_with(
+            crate::harness::LocalRegistry::new(),
+            test_session_id(),
+            Default::default(),
+        );
+        let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id());
+        bridge
+            .emit(SessionEvent::PhaseChanged {
+                phase: SessionPhase::ToolExecution,
+            })
+            .await;
+    }
+
+    #[tokio::test]
+    async fn emit_all_event_variants_does_not_panic() {
+        let harness = ToolHarness::local_only_with(
+            crate::harness::LocalRegistry::new(),
+            test_session_id(),
+            Default::default(),
+        );
+        let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id());
+
+        // Smoke test: every variant emits without panic through a
+        // local-only harness. Includes error/cancelled outcomes to
+        // cover non-happy-path enum values.
+        let events = vec![
+            SessionEvent::TurnStarted {
+                turn_number: 1,
+                model_id: "grok-3".into(),
+                yolo_mode: false,
+            },
+            SessionEvent::ToolCallStarted {
+                tool_call_id: "c1".into(),
+                tool_name: "bash".into(),
+                turn_number: 1,
+            },
+            SessionEvent::ToolCallCompleted {
+                tool_call_id: "c1".into(),
+                tool_name: "bash".into(),
+                duration_ms: 100,
+                outcome: ToolCallOutcome::Success,
+            },
+            SessionEvent::ToolCallCompleted {
+                tool_call_id: "c2".into(),
+                tool_name: "read_file".into(),
+                duration_ms: 50,
+                outcome: ToolCallOutcome::Error,
+            },
+            SessionEvent::ToolCallCompleted {
+                tool_call_id: "c3".into(),
+                tool_name: "grep".into(),
+                duration_ms: 10,
+                outcome: ToolCallOutcome::Cancelled,
+            },
+            SessionEvent::PhaseChanged {
+                phase: SessionPhase::Idle,
+            },
+            SessionEvent::TurnEnded {
+                turn_number: 1,
+                outcome: TurnHookOutcome::Completed,
+                duration_ms: 500,
+                tool_call_count: 3,
+                model_id: "grok-3".into(),
+            },
+            SessionEvent::TurnEnded {
+                turn_number: 2,
+                outcome: TurnHookOutcome::Error,
+                duration_ms: 100,
+                tool_call_count: 0,
+                model_id: "grok-3".into(),
+            },
+            SessionEvent::Unknown,
+        ];
+
+        for event in events {
+            bridge.emit(event).await;
+        }
+    }
+}

crates/common/xai-computer-hub-sdk/src/oidc_provider.rs

@@ -0,0 +1,336 @@
+//! [`AuthProvider`] that refreshes OIDC tokens before they expire.
+//!
+//! `current()` checks token expiry and, if needed, performs OIDC
+//! discovery + token exchange before returning the credential.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use chrono::{DateTime, Utc};
+use parking_lot::Mutex;
+
+use crate::auth::{AuthCredential, AuthIdentity, AuthProvider};
+
+pub type OnRefreshCallback = Arc<dyn Fn(&RefreshEvent) + Send + Sync>;
+
+#[derive(Debug, Clone)]
+pub struct RefreshEvent {
+    pub access_token: String,
+    pub new_refresh_token: Option<String>,
+    pub expires_at: Option<DateTime<Utc>>,
+}
+
+struct TokenState {
+    access_token: String,
+    refresh_token: String,
+    expires_at: Option<DateTime<Utc>>,
+}
+
+pub struct OidcAuthProvider {
+    state: Mutex<TokenState>,
+    issuer: String,
+    client_id: String,
+    user_id: Option<String>,
+    principal_type: Option<String>,
+    principal_id: Option<String>,
+    on_refresh: Option<OnRefreshCallback>,
+}
+
+const REFRESH_MARGIN: Duration = Duration::from_secs(60);
+
+impl std::fmt::Debug for OidcAuthProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("OidcAuthProvider")
+            .field("issuer", &self.issuer)
+            .field("client_id", &self.client_id)
+            .finish_non_exhaustive()
+    }
+}
+
+pub struct OidcAuthProviderBuilder {
+    access_token: String,
+    refresh_token: String,
+    issuer: String,
+    client_id: String,
+    expires_at: Option<DateTime<Utc>>,
+    user_id: Option<String>,
+    principal_type: Option<String>,
+    principal_id: Option<String>,
+    on_refresh: Option<OnRefreshCallback>,
+}
+
+impl OidcAuthProviderBuilder {
+    pub fn new(
+        access_token: impl Into<String>,
+        refresh_token: impl Into<String>,
+        issuer: impl Into<String>,
+        client_id: impl Into<String>,
+    ) -> Self {
+        Self {
+            access_token: access_token.into(),
+            refresh_token: refresh_token.into(),
+            issuer: issuer.into(),
+            client_id: client_id.into(),
+            expires_at: None,
+            user_id: None,
+            principal_type: None,
+            principal_id: None,
+            on_refresh: None,
+        }
+    }
+
+    pub fn expires_at(mut self, expires_at: DateTime<Utc>) -> Self {
+        self.expires_at = Some(expires_at);
+        self
+    }
+
+    /// Owner user id parsed from the auth source, surfaced via
+    /// [`AuthProvider::identity`].
+    pub fn user_id(mut self, user_id: impl Into<String>) -> Self {
+        self.user_id = Some(user_id.into());
+        self
+    }
+
+    pub fn principal_type(mut self, pt: impl Into<String>) -> Self {
+        self.principal_type = Some(pt.into());
+        self
+    }
+
+    pub fn principal_id(mut self, pid: impl Into<String>) -> Self {
+        self.principal_id = Some(pid.into());
+        self
+    }
+
+    pub fn on_refresh(mut self, cb: OnRefreshCallback) -> Self {
+        self.on_refresh = Some(cb);
+        self
+    }
+
+    pub fn build(self) -> OidcAuthProvider {
+        OidcAuthProvider {
+            state: Mutex::new(TokenState {
+                access_token: self.access_token,
+                refresh_token: self.refresh_token,
+                expires_at: self.expires_at,
+            }),
+            issuer: self.issuer,
+            client_id: self.client_id,
+            user_id: self.user_id,
+            principal_type: self.principal_type,
+            principal_id: self.principal_id,
+            on_refresh: self.on_refresh,
+        }
+    }
+}
+
+impl AuthProvider for OidcAuthProvider {
+    fn current(&self) -> AuthCredential {
+        let expired = {
+            let s = self.state.lock();
+            s.expires_at.is_some_and(|exp| {
+                Utc::now() + chrono::Duration::from_std(REFRESH_MARGIN).unwrap() >= exp
+            })
+        };
+        if expired && let Err(e) = self.try_refresh() {
+            tracing::warn!(error = %e, "OIDC refresh failed, using stale token");
+        }
+        let s = self.state.lock();
+        AuthCredential::bearer(&s.access_token)
+    }
+
+    /// Surface the principal fields parsed from the auth source. `None` only
+    /// when no `user_id` was supplied (nothing to attribute).
+    fn identity(&self) -> Option<AuthIdentity> {
+        let user_id = self.user_id.clone()?;
+        Some(AuthIdentity {
+            user_id,
+            principal_type: self.principal_type.clone(),
+            principal_id: self.principal_id.clone(),
+        })
+    }
+}
+
+impl OidcAuthProvider {
+    fn try_refresh(&self) -> Result<(), Box<dyn std::error::Error>> {
+        tracing::info!(issuer = %self.issuer, "refreshing OIDC token");
+        if let Ok(handle) = tokio::runtime::Handle::try_current() {
+            tokio::task::block_in_place(|| handle.block_on(self.do_refresh()))
+        } else {
+            tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()?
+                .block_on(self.do_refresh())
+        }
+    }
+
+    async fn do_refresh(&self) -> Result<(), Box<dyn std::error::Error>> {
+        let refresh_token = self.state.lock().refresh_token.clone();
+        let issuer = self.issuer.trim_end_matches('/');
+        let client = reqwest::Client::new();
+
+        #[derive(serde::Deserialize)]
+        struct Discovery {
+            token_endpoint: String,
+        }
+
+        let disc: Discovery = client
+            .get(format!("{issuer}/.well-known/openid-configuration"))
+            .timeout(Duration::from_secs(10))
+            .send()
+            .await?
+            .error_for_status()?
+            .json()
+            .await?;
+
+        let mut params = vec![
+            ("grant_type", "refresh_token"),
+            ("refresh_token", refresh_token.as_str()),
+            ("client_id", self.client_id.as_str()),
+        ];
+        let pt = self.principal_type.clone();
+        let pid = self.principal_id.clone();
+        if let Some(ref v) = pt {
+            params.push(("principal_type", v));
+        }
+        if let Some(ref v) = pid {
+            params.push(("principal_id", v));
+        }
+
+        #[derive(serde::Deserialize)]
+        struct Tokens {
+            access_token: String,
+            #[serde(default)]
+            refresh_token: Option<String>,
+            #[serde(default)]
+            expires_in: Option<u64>,
+        }
+
+        let tokens: Tokens = client
+            .post(&disc.token_endpoint)
+            .form(&params)
+            .timeout(Duration::from_secs(15))
+            .send()
+            .await?
+            .error_for_status()?
+            .json()
+            .await?;
+
+        let expires_at = tokens
+            .expires_in
+            .map(|s| Utc::now() + chrono::Duration::seconds(s as i64));
+
+        tracing::info!(expires_at = ?expires_at, "OIDC token refreshed");
+
+        if let Some(ref cb) = self.on_refresh {
+            cb(&RefreshEvent {
+                access_token: tokens.access_token.clone(),
+                new_refresh_token: tokens.refresh_token.clone(),
+                expires_at,
+            });
+        }
+
+        let mut s = self.state.lock();
+        s.access_token = tokens.access_token;
+        if let Some(rt) = tokens.refresh_token {
+            s.refresh_token = rt;
+        }
+        s.expires_at = expires_at;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn current_returns_token_when_not_expired() {
+        let provider = OidcAuthProviderBuilder::new(
+            "access-tok",
+            "refresh-tok",
+            "https://auth.example.com",
+            "client1",
+        )
+        .expires_at(Utc::now() + chrono::Duration::hours(1))
+        .build();
+
+        let cred = provider.current();
+        match cred {
+            AuthCredential::Bearer { token } => {
+                assert_eq!(token, "access-tok");
+            }
+            _ => panic!("expected Bearer"),
+        }
+    }
+
+    #[test]
+    fn current_returns_token_when_no_expiry() {
+        let provider = OidcAuthProviderBuilder::new(
+            "no-expiry-tok",
+            "refresh-tok",
+            "https://auth.example.com",
+            "client1",
+        )
+        .build();
+
+        let cred = provider.current();
+        match cred {
+            AuthCredential::Bearer { token } => assert_eq!(token, "no-expiry-tok"),
+            _ => panic!("expected Bearer"),
+        }
+    }
+
+    #[test]
+    fn current_returns_stale_token_when_refresh_fails() {
+        // Expired token, but issuer is unreachable — should return stale
+        let provider = OidcAuthProviderBuilder::new(
+            "stale-tok",
+            "refresh-tok",
+            "https://localhost:1", // unreachable
+            "client1",
+        )
+        .expires_at(Utc::now() - chrono::Duration::hours(1))
+        .build();
+
+        let cred = provider.current();
+        match cred {
+            AuthCredential::Bearer { token } => assert_eq!(token, "stale-tok"),
+            _ => panic!("expected Bearer"),
+        }
+    }
+
+    #[test]
+    fn identity_surfaces_principal_fields() {
+        let provider = OidcAuthProviderBuilder::new("tok", "rt", "https://auth.example.com", "c1")
+            .user_id("user-1")
+            .principal_type("Team")
+            .principal_id("team-9")
+            .build();
+        let id = provider.identity().expect("identity present");
+        assert_eq!(id.user_id, "user-1");
+        assert_eq!(id.principal_type.as_deref(), Some("Team"));
+        assert_eq!(id.principal_id.as_deref(), Some("team-9"));
+    }
+
+    #[test]
+    fn identity_none_without_user_id() {
+        let provider =
+            OidcAuthProviderBuilder::new("tok", "rt", "https://auth.example.com", "c1").build();
+        assert!(provider.identity().is_none());
+    }
+
+    #[test]
+    fn debug_does_not_leak_tokens() {
+        let provider = OidcAuthProviderBuilder::new(
+            "secret-access-token",
+            "secret-refresh-token",
+            "https://auth.example.com",
+            "client1",
+        )
+        .build();
+
+        let debug = format!("{provider:?}");
+        assert!(!debug.contains("secret-access-token"));
+        assert!(!debug.contains("secret-refresh-token"));
+    }
+}

crates/common/xai-computer-hub-sdk/src/pool.rs

@@ -0,0 +1,344 @@
+//! Process-wide connection pool keyed by `(url, principal)`.
+//!
+//! Two [`crate::ToolServer`] builds with the same `(url, credential)`
+//! observe the same `Arc<HubConnection>`; distinct credentials open
+//! distinct sockets. The pool is the canonical entry point — direct
+//! [`crate::HubConnection::connect`] calls are reserved for tests and
+//! one-shot programs that explicitly want unpooled behaviour.
+
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use dashmap::DashMap;
+use tokio::sync::OnceCell;
+use tokio::task::JoinHandle;
+use url::Url;
+use xai_tool_protocol::ConnectionKind;
+
+use crate::auth::AuthProvider;
+use crate::connection::{
+    ConnKey, ConnectCallback, ConnectionConfig, ConnectionTuning, DisconnectCallback,
+    HubConnection, ReconnectCallback,
+};
+use crate::error::ClientError;
+
+/// Idle window for the reaper: a pooled connection is evictable once it is
+/// unused (`Arc::strong_count == 1`, i.e. only the pool holds it) **and**
+/// `now - last_handout >= DEFAULT_POOL_IDLE_TTL`.
+///
+/// Note the clock is `last_handout` (the last time the pool returned the
+/// connection), not the moment the last consumer `Arc` was dropped: a
+/// connection held longer than the TTL and then released is eligible on the
+/// very next sweep, with no extra post-drop grace period. The only hard
+/// guarantee is that an in-use connection (`strong_count > 1`) is never
+/// reaped. Tuned well above the server's own 90s dead-peer idle timeout so a
+/// short borrow between turns of an active conversation isn't churned.
+pub const DEFAULT_POOL_IDLE_TTL: Duration = Duration::from_secs(300);
+
+/// How often the shared pool's idle reaper scans for evictable entries.
+pub const DEFAULT_POOL_SWEEP_INTERVAL: Duration = Duration::from_secs(60);
+
+/// A pooled connection plus the last time it was handed out to a caller.
+///
+/// `last_handout` is refreshed on every [`HubConnectionPool::get_or_connect`]
+/// hit (and on the initial insert), so a connection that is repeatedly
+/// re-fetched never looks idle even if its [`Arc`] strong count briefly
+/// returns to 1 between fetches. Eviction additionally requires
+/// `Arc::strong_count == 1` (only the pool holds it), so a connection a
+/// consumer still holds is never reaped regardless of `last_handout`.
+struct Pooled {
+    conn: Arc<HubConnection>,
+    last_handout: Instant,
+}
+
+/// The process-global pool used by [`HubConnectionPool::shared`].
+///
+/// `tokio::sync::OnceCell` is preferred over `std::sync::OnceLock` /
+/// `LazyLock` here because the pool is only ever observed from
+/// async contexts (the connection actor lives on a tokio runtime
+/// already), so the async-aware `get_or_init` semantics avoid the
+/// blocking-init footgun of the sync alternatives without taking a
+/// hard dependency on additional sync primitives.
+///
+/// Tests MUST use [`HubConnectionPool::new`] to avoid cross-test
+/// pollution: cargo runs all integration tests in the same binary
+/// unless otherwise configured, so any test that touches
+/// `HubConnectionPool::shared()` leaves the pool populated for
+/// subsequent tests.
+static SHARED: OnceCell<Arc<HubConnectionPool>> = OnceCell::const_new();
+
+/// Pool of live server connections.
+pub struct HubConnectionPool {
+    connections: DashMap<ConnKey, Pooled>,
+}
+
+impl std::fmt::Debug for HubConnectionPool {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("HubConnectionPool")
+            .field("connection_count", &self.connections.len())
+            .finish()
+    }
+}
+
+impl HubConnectionPool {
+    /// Build a fresh, unshared pool. Tests typically use this so each
+    /// test sees an isolated registry.
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {
+            connections: DashMap::new(),
+        })
+    }
+
+    /// Return the process-wide shared pool, lazily initialising it on
+    /// the first call. Subsequent callers in the same process observe
+    /// the same `Arc`.
+    ///
+    /// The shared pool spawns an idle reaper (see [`Self::spawn_idle_reaper`])
+    /// exactly once, so a connection that is unused (`strong_count == 1`) and
+    /// has not been handed out for [`DEFAULT_POOL_IDLE_TTL`] is closed instead
+    /// of living for the whole process lifetime. (Unpooled / test pools built
+    /// via [`Self::new`] do not
+    /// get a reaper; they can call [`Self::sweep_idle`] directly.)
+    pub async fn shared() -> Arc<Self> {
+        SHARED
+            .get_or_init(|| async {
+                let pool = Self::new();
+                pool.spawn_idle_reaper(DEFAULT_POOL_IDLE_TTL, DEFAULT_POOL_SWEEP_INTERVAL);
+                pool
+            })
+            .await
+            .clone()
+    }
+
+    /// Look up an existing pooled connection for `(url, credential)`,
+    /// or open a fresh one if no pooled entry exists.
+    ///
+    /// `kind` is the connection role announced in the hello frame. The
+    /// pool is keyed by `(url, principal)` only; mixing
+    /// [`ConnectionKind`] values for the same `(url, principal)` is a
+    /// caller error and surfaces as a [`ClientError::InvalidConfig`].
+    ///
+    /// The optional extra access key is not part of the pool key, so the first
+    /// caller's key is the one carried on a shared connection's handshake (in
+    /// practice it is a per-deployment constant). The plaintext-scheme guard is
+    /// re-checked on every call below so it can't be bypassed by a cached
+    /// insecure entry.
+    pub async fn get_or_connect(
+        self: &Arc<Self>,
+        url: Url,
+        credential: Arc<dyn AuthProvider>,
+        kind: ConnectionKind,
+        on_reconnect: Option<Arc<ReconnectCallback>>,
+        on_disconnect: Option<Arc<DisconnectCallback>>,
+        server_id: Option<xai_tool_protocol::ServerId>,
+        alpha_test_key: Option<String>,
+        allow_insecure_ws: bool,
+    ) -> Result<Arc<HubConnection>, ClientError> {
+        self.get_or_connect_tuned(
+            url,
+            credential,
+            kind,
+            on_reconnect,
+            on_disconnect,
+            None, // on_connect (unused by the simple wrapper)
+            server_id,
+            None,
+            None,
+            alpha_test_key,
+            allow_insecure_ws,
+            ConnectionTuning::default(),
+        )
+        .await
+    }
+
+    /// Like [`Self::get_or_connect`] but carries optional connection-tuning
+    /// knobs ([`ConnectionTuning`]) onto a freshly-opened connection. A
+    /// `Default` tuning is behaviourally identical to `get_or_connect`, so
+    /// existing callers are unaffected.
+    ///
+    /// Tuning binds to the socket at open time: it takes effect only when
+    /// THIS call opens the connection. Because the pool dedups by
+    /// `(url, principal)`, a hit on an existing entry returns that
+    /// connection as-is and the `tuning` argument is ignored — the first
+    /// opener's ping/backoff settings win for the lifetime of the pooled
+    /// connection. Callers that need distinct tuning must use a distinct
+    /// `(url, principal)` or an unpooled [`HubConnection::connect`].
+    pub(crate) async fn get_or_connect_tuned(
+        self: &Arc<Self>,
+        url: Url,
+        credential: Arc<dyn AuthProvider>,
+        kind: ConnectionKind,
+        on_reconnect: Option<Arc<ReconnectCallback>>,
+        on_disconnect: Option<Arc<DisconnectCallback>>,
+        on_connect: Option<Arc<ConnectCallback>>,
+        server_id: Option<xai_tool_protocol::ServerId>,
+        server_description: Option<String>,
+        server_metadata: Option<serde_json::Value>,
+        alpha_test_key: Option<String>,
+        allow_insecure_ws: bool,
+        tuning: ConnectionTuning,
+    ) -> Result<Arc<HubConnection>, ClientError> {
+        if url.scheme() != "wss" && !crate::connection::host_is_loopback(&url) && !allow_insecure_ws
+        {
+            return Err(ClientError::InsecureScheme { url });
+        }
+        let key = ConnKey {
+            url: url.as_str().to_owned(),
+            principal: credential.principal_key(),
+        };
+        if let Some(mut existing) = self.connections.get_mut(&key) {
+            existing.last_handout = Instant::now();
+            let conn = existing.conn.clone();
+            drop(existing);
+            if conn.kind() != kind {
+                return Err(ClientError::InvalidConfig(format!(
+                    "pool entry for {} bound to {:?}; rebuild requested {:?}",
+                    key.url,
+                    conn.kind(),
+                    kind
+                )));
+            }
+            return Ok(conn);
+        }
+        let config = ConnectionConfig {
+            url,
+            credential,
+            kind,
+            on_reconnect,
+            on_disconnect,
+            on_connect,
+            server_id,
+            server_description,
+            server_metadata,
+            outbound_buffer: None,
+            tuning,
+            alpha_test_key,
+            allow_insecure_ws,
+            on_fatal: Some(Arc::downgrade(self)),
+        };
+        let conn = HubConnection::connect(config).await?;
+        // Race window: another caller may have inserted between our
+        // `get` and `connect`. Resolve via `entry().or_insert_with`
+        // semantics — if we lose the race we drop our fresh
+        // connection and adopt the winning one.
+        match self.connections.entry(key.clone()) {
+            dashmap::Entry::Occupied(mut existing) => {
+                existing.get_mut().last_handout = Instant::now();
+                let winner = existing.get().conn.clone();
+                drop(conn);
+                if winner.kind() != kind {
+                    return Err(ClientError::InvalidConfig(format!(
+                        "pool entry for {} bound to {:?}; rebuild requested {:?}",
+                        key.url,
+                        winner.kind(),
+                        kind
+                    )));
+                }
+                Ok(winner)
+            }
+            dashmap::Entry::Vacant(slot) => {
+                crate::metrics::pool_connections_inc();
+                slot.insert(Pooled {
+                    conn: conn.clone(),
+                    last_handout: Instant::now(),
+                });
+                Ok(conn)
+            }
+        }
+    }
+
+    /// Number of pooled connections. Intended for tests and metrics.
+    pub fn len(&self) -> usize {
+        self.connections.len()
+    }
+
+    /// `true` when no connection is pooled.
+    pub fn is_empty(&self) -> bool {
+        self.connections.is_empty()
+    }
+
+    /// Forget the pooled connection for `key`. The actual underlying
+    /// `Arc<HubConnection>` is dropped only when no other holder
+    /// keeps a reference; the next [`Self::get_or_connect`] for the
+    /// same key opens a fresh socket.
+    pub fn forget(&self, key: &ConnKey) {
+        if self.connections.remove(key).is_some() {
+            crate::metrics::pool_connections_dec();
+        }
+    }
+
+    /// Close and remove every pooled connection that is BOTH unused (no live
+    /// consumer holds an `Arc` — only the pool does, so `strong_count == 1`)
+    /// AND idle longer than `idle_ttl` (no hand-out within the window).
+    /// Removing the entry drops the pool's last `Arc<HubConnection>`, whose
+    /// `Drop` closes the socket.
+    ///
+    /// The strong-count check runs inside the map's per-shard lock (via
+    /// [`DashMap::retain`]), serialised against `get_or_connect`, so a
+    /// connection handed out concurrently is never evicted out from under a
+    /// caller. Returns the number of connections evicted.
+    pub fn sweep_idle(&self, idle_ttl: Duration) -> usize {
+        let now = Instant::now();
+        let mut evicted = 0usize;
+        self.connections.retain(|_key, pooled| {
+            let idle_for = now.saturating_duration_since(pooled.last_handout);
+            // `strong_count == 1` ⇒ only this pool entry references the
+            // connection, so no consumer can still be using it.
+            let unused = Arc::strong_count(&pooled.conn) == 1;
+            let evict = unused && idle_for >= idle_ttl;
+            if evict {
+                evicted += 1;
+            }
+            !evict
+        });
+        for _ in 0..evicted {
+            crate::metrics::pool_connections_dec();
+            crate::metrics::pool_evictions_inc();
+        }
+        evicted
+    }
+
+    /// Spawn a background task that calls [`Self::sweep_idle`] every
+    /// `sweep_interval`, closing connections idle longer than `idle_ttl`.
+    ///
+    /// The task holds a [`std::sync::Weak`] to the pool, so it exits on its
+    /// own once the last strong `Arc<HubConnectionPool>` is dropped (it never
+    /// keeps the pool alive). The first interval tick is skipped so a
+    /// freshly-handed-out connection is never swept on the immediate tick.
+    pub fn spawn_idle_reaper(
+        self: &Arc<Self>,
+        idle_ttl: Duration,
+        sweep_interval: Duration,
+    ) -> JoinHandle<()> {
+        let weak = Arc::downgrade(self);
+        tokio::spawn(async move {
+            let mut ticker = tokio::time::interval(sweep_interval);
+            ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+            // `interval`'s first tick resolves immediately; skip it.
+            ticker.tick().await;
+            loop {
+                ticker.tick().await;
+                let Some(pool) = weak.upgrade() else { break };
+                pool.sweep_idle(idle_ttl);
+            }
+        })
+    }
+
+    /// Like [`Self::forget`] but identity-checked: only removes the slot
+    /// when `predicate` accepts the currently-stored connection. The
+    /// self-evicting actor passes an `Arc::ptr_eq` check so a race-loser
+    /// can never drop the winner's fresh entry (ABA-safe).
+    pub(crate) fn forget_if(
+        &self,
+        key: &ConnKey,
+        predicate: impl FnOnce(&Arc<HubConnection>) -> bool,
+    ) {
+        if self
+            .connections
+            .remove_if(key, |_, pooled| predicate(&pooled.conn))
+            .is_some()
+        {
+            crate::metrics::pool_connections_dec();
+        }
+    }
+}

crates/common/xai-computer-hub-sdk/src/refcount.rs

@@ -0,0 +1,123 @@
+//! Generic refcounted-binding helper used by the connection's
+//! bound-session set.
+//!
+//! Multiple [`crate::ToolServer`] instances can share one
+//! [`crate::HubConnection`] when they target the same `(url, principal)`.
+//! Each instance independently asks for a session binding; the substrate
+//! must `register_session` once per session (not once per consumer) and
+//! `unregister_session` only when the LAST consumer drops its borrow.
+//! [`RefCountedSet`] tracks the per-key borrow count behind a
+//! [`dashmap::DashMap`] so increments and decrements never serialise on
+//! a single mutex.
+
+use std::hash::Hash;
+
+use dashmap::DashMap;
+
+/// Refcounted set keyed by `K`. Each [`Self::increment`] returns the
+/// new count; the corresponding [`Self::decrement`] returns the count
+/// AFTER the decrement (so callers fire teardown when the result is
+/// `Some(0)`).
+#[derive(Debug, Default)]
+pub struct RefCountedSet<K: Eq + Hash> {
+    counts: DashMap<K, u64>,
+}
+
+impl<K: Eq + Hash> RefCountedSet<K> {
+    /// Empty set.
+    pub fn new() -> Self {
+        Self {
+            counts: DashMap::new(),
+        }
+    }
+
+    /// Increment `key`'s refcount. Returns `(prev_count, new_count)`
+    /// so callers can detect the 0→1 edge (when the protocol-level
+    /// register call must fire).
+    pub fn increment(&self, key: K) -> (u64, u64)
+    where
+        K: Clone,
+    {
+        let mut entry = self.counts.entry(key).or_insert(0);
+        let prev = *entry;
+        *entry = prev.saturating_add(1);
+        (prev, *entry)
+    }
+
+    /// Decrement `key`'s refcount. Returns the post-decrement count;
+    /// `Some(0)` means the entry was removed and callers should fire
+    /// the protocol-level unregister. `None` means the key was not
+    /// present (idempotent drop).
+    pub fn decrement(&self, key: &K) -> Option<u64> {
+        let mut current = None;
+        self.counts.remove_if_mut(key, |_, value| {
+            *value = value.saturating_sub(1);
+            current = Some(*value);
+            *value == 0
+        });
+        current
+    }
+
+    /// Snapshot the live keys. Allocates a fresh `Vec` — only used by
+    /// the reconnect-replay path which fires once per disconnect.
+    pub fn snapshot_keys(&self) -> Vec<K>
+    where
+        K: Clone,
+    {
+        self.counts.iter().map(|kv| kv.key().clone()).collect()
+    }
+
+    /// `true` when no key has a non-zero refcount.
+    pub fn is_empty(&self) -> bool {
+        self.counts.is_empty()
+    }
+
+    /// Number of distinct live keys.
+    pub fn len(&self) -> usize {
+        self.counts.len()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn increment_returns_new_count() {
+        let set = RefCountedSet::<&'static str>::new();
+        assert_eq!(set.increment("a"), (0, 1));
+        assert_eq!(set.increment("a"), (1, 2));
+        assert_eq!(set.increment("b"), (0, 1));
+        assert_eq!(set.len(), 2);
+    }
+
+    #[test]
+    fn decrement_removes_at_zero() {
+        let set = RefCountedSet::<&'static str>::new();
+        set.increment("a");
+        set.increment("a");
+        assert_eq!(set.decrement(&"a"), Some(1));
+        assert!(!set.is_empty());
+        assert_eq!(set.decrement(&"a"), Some(0));
+        assert!(set.is_empty());
+    }
+
+    #[test]
+    fn decrement_unknown_returns_none() {
+        let set = RefCountedSet::<&'static str>::new();
+        assert!(set.decrement(&"missing").is_none());
+    }
+
+    #[test]
+    fn increment_saturates_at_u64_max() {
+        let set = RefCountedSet::<&'static str>::new();
+        // Pre-load the entry to MAX-1 via direct DashMap access. The
+        // public API only ever reaches this region via overflow,
+        // which is impossible in practice; this test pins the
+        // saturating_add defensive line so it can't silently regress
+        // to wrapping_add.
+        set.counts.insert("max", u64::MAX - 1);
+        assert_eq!(set.increment("max"), (u64::MAX - 1, u64::MAX));
+        assert_eq!(set.increment("max"), (u64::MAX, u64::MAX));
+    }
+}

crates/common/xai-computer-hub-sdk/src/trace_donate.rs

@@ -0,0 +1,206 @@
+//! Forward selected spans to the connected server over the WebSocket
+//! transport (`traces.donate`). The bounded retry buffer + drain barrier
+//! live in [`crate::donate_pump`]; overflow drops spans — telemetry,
+//! never correctness.
+
+use std::borrow::Cow;
+
+use base64::Engine as _;
+use fastrace::collector::{Reporter, SpanRecord};
+use fastrace_opentelemetry::OpenTelemetryReporter;
+use opentelemetry::InstrumentationScope;
+use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
+use opentelemetry_proto::transform::common::tonic::ResourceAttributesWithSchema;
+use opentelemetry_proto::transform::trace::tonic::group_spans_by_resource_and_scope;
+use opentelemetry_sdk::Resource;
+use opentelemetry_sdk::error::OTelSdkResult;
+use opentelemetry_sdk::trace::{SpanData, SpanExporter};
+use prost::Message as _;
+use tokio::sync::mpsc;
+use xai_tool_protocol::{MAX_DONATION_BYTES, MAX_SPANS_PER_DONATION};
+
+use crate::donate_pump::{PENDING_FLUSHES, PumpMsg, drain_via, run_pump};
+use crate::server::ToolServer;
+
+/// fastrace [`Reporter`] feeding the donation pump.
+pub struct HubDonatingReporter(OpenTelemetryReporter);
+
+impl Reporter for HubDonatingReporter {
+    fn report(&mut self, spans: Vec<SpanRecord>) {
+        if spans.is_empty() {
+            return;
+        }
+        self.0.report(spans);
+    }
+}
+
+/// [`SpanExporter`] that encodes OTLP requests onto the pump channel.
+/// Runs on fastrace's collector thread; must never block.
+#[derive(Debug)]
+struct PumpSpanExporter {
+    tx: mpsc::Sender<PumpMsg>,
+    resource: ResourceAttributesWithSchema,
+}
+
+impl SpanExporter for PumpSpanExporter {
+    fn export(
+        &self,
+        batch: Vec<SpanData>,
+    ) -> impl std::future::Future<Output = OTelSdkResult> + Send {
+        let mut remaining = batch;
+        while !remaining.is_empty() {
+            let chunk = if remaining.len() > MAX_SPANS_PER_DONATION {
+                let rest = remaining.split_off(MAX_SPANS_PER_DONATION);
+                std::mem::replace(&mut remaining, rest)
+            } else {
+                std::mem::take(&mut remaining)
+            };
+            let request = ExportTraceServiceRequest {
+                resource_spans: group_spans_by_resource_and_scope(chunk, &self.resource),
+            };
+            let bytes = request.encode_to_vec();
+            if bytes.len() > MAX_DONATION_BYTES {
+                tracing::debug!(len = bytes.len(), "dropping oversized donation payload");
+                continue;
+            }
+            let payload = base64::engine::general_purpose::STANDARD.encode(bytes);
+            if self.tx.try_send(PumpMsg::Payload(payload)).is_err() {
+                tracing::debug!("trace donation queue full; dropping span batch");
+            }
+        }
+        std::future::ready(Ok(()))
+    }
+
+    fn set_resource(&mut self, resource: &Resource) {
+        self.resource = resource.into();
+    }
+}
+
+/// Shutdown fence: drains queued donations before the connection closes.
+pub struct TraceDonationPump {
+    tx: mpsc::Sender<PumpMsg>,
+}
+
+impl TraceDonationPump {
+    /// Resolves once every payload queued before this call has had a
+    /// send attempt. Call after `fastrace::flush()`.
+    pub async fn drain(&self) {
+        drain_via(&self.tx).await;
+    }
+}
+
+impl ToolServer {
+    /// Spawn the donation pump and return its reporter + drain handle.
+    /// `service_name` must be server-allowlisted.
+    pub fn trace_donation_reporter(
+        &self,
+        service_name: impl Into<String>,
+    ) -> (HubDonatingReporter, TraceDonationPump) {
+        let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
+        let server = self.downgrade();
+        tokio::spawn(run_pump(rx, move |payload: String| {
+            let server = server.clone();
+            async move {
+                let Some(server) = server.upgrade() else {
+                    return (false, payload);
+                };
+                let ok = server.donate_traces(&payload).await.is_ok();
+                (ok, payload)
+            }
+        }));
+        self.set_donation_pump(tx.clone());
+
+        let resource = Resource::builder()
+            .with_service_name(service_name.into())
+            .build();
+        let exporter = PumpSpanExporter {
+            tx: tx.clone(),
+            resource: (&resource).into(),
+        };
+        let reporter = OpenTelemetryReporter::new(
+            exporter,
+            Cow::Owned(resource),
+            InstrumentationScope::default(),
+        );
+        (HubDonatingReporter(reporter), TraceDonationPump { tx })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::SystemTime;
+
+    use opentelemetry::trace::{SpanContext, SpanKind, Status, TraceFlags, TraceState};
+
+    use super::*;
+
+    #[tokio::test]
+    async fn exporter_encodes_standard_otlp_with_resource() {
+        let resource = Resource::builder()
+            .with_service_name("test-service")
+            .build();
+        let (tx, mut rx) = mpsc::channel::<PumpMsg>(4);
+        let exporter = PumpSpanExporter {
+            tx,
+            resource: (&resource).into(),
+        };
+
+        let span = SpanData {
+            span_context: SpanContext::new(
+                0x0af7651916cd43dd8448eb211c80319c_u128.into(),
+                0xb7ad6b7169203331_u64.into(),
+                TraceFlags::SAMPLED,
+                false,
+                TraceState::default(),
+            ),
+            parent_span_id: 0_u64.into(),
+            parent_span_is_remote: false,
+            span_kind: SpanKind::Internal,
+            name: "tool_server.tool_call".into(),
+            start_time: SystemTime::UNIX_EPOCH,
+            end_time: SystemTime::UNIX_EPOCH,
+            attributes: vec![opentelemetry::KeyValue::new("tool_id", "bash")],
+            dropped_attributes_count: 0,
+            events: opentelemetry_sdk::trace::SpanEvents::default(),
+            links: opentelemetry_sdk::trace::SpanLinks::default(),
+            status: Status::Unset,
+            instrumentation_scope: InstrumentationScope::default(),
+        };
+        exporter
+            .export(vec![span])
+            .await
+            .expect("export must succeed");
+
+        let Some(PumpMsg::Payload(payload)) = rx.try_recv().ok() else {
+            panic!("exporter must enqueue one payload");
+        };
+        let bytes = base64::engine::general_purpose::STANDARD
+            .decode(payload)
+            .expect("payload must be base64");
+        let request =
+            ExportTraceServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP");
+        let resource_spans = &request.resource_spans[0];
+        let service_name = resource_spans
+            .resource
+            .as_ref()
+            .unwrap()
+            .attributes
+            .iter()
+            .find(|kv| kv.key == "service.name")
+            .and_then(|kv| kv.value.as_ref())
+            .map(|v| format!("{v:?}"));
+        assert!(
+            service_name.unwrap_or_default().contains("test-service"),
+            "resource must carry the donor service.name"
+        );
+        let span = &resource_spans.scope_spans[0].spans[0];
+        assert_eq!(span.name, "tool_server.tool_call");
+        assert_eq!(
+            format!(
+                "{:032x}",
+                u128::from_be_bytes(span.trace_id.as_slice().try_into().unwrap())
+            ),
+            "0af7651916cd43dd8448eb211c80319c"
+        );
+    }
+}