salience: add gRPC client + TLS plumbing for stateful vllm sessions
Adds the client-side of a stateful gRPC protocol against vllm, plus the TLS trust machinery so we can talk to self-signed vllm servers. Protocol (proto/salience.proto): Bidi-streaming Session RPC carries OpenSession / AppendTokens / Generate / Cancel from client and SessionReady / PrefillProgress / Token / GenerateDone / Error from server. Separate Fork unary RPC for cheap branching (prefix cache shares KV automatically). Plus ListSessions, CloseSession, GetReadoutManifest admin RPCs. Per-token readouts ship as packed f32 ([n_layers * n_concepts] per token, flat). Logprobs use range-selected positions plus a top-k parameter — empty ranges means no logprobs, any range means emit sampled-token logprob at those positions, top_k > 0 adds alternatives. Client (src/agent/api/salience.rs): Tonic-generated types under pb::, a connect() helper, with_auth() for bearer metadata, and a Session handle wrapping the bidi stream: open() handshakes SessionReady; append() is fire-and-forget; generate() returns impl Stream<Item = Event> that drains inbound until Done or terminating Error. One generate at a time per session. Peak picker (src/agent/salience.rs): Pure function over ReadoutEntry traces. Per-concept z-score against trace global stats; contiguous above-threshold regions emit one peak at the local max. Configurable sigma threshold and min-std safety floor. Deterministic tie-break on offset then concept name. 12 unit tests covering empty traces, flat channels, single/multi spikes, contiguous humps, multi-concept independence, trailing runs, sub-threshold noise, layer-out-of-range, manifest shape mismatch, and threshold tunability. TLS (src/agent/api/http.rs): HttpClient::build now also loads every .pem file under ~/.consciousness/certs/ into the rustls root store — so dropping a <host>.pem in that directory is enough to trust a new self- signed server; no code changes per new host. Also installs the rustls default crypto provider explicitly via OnceLock: tonic's tls features pulled in both ring and aws-lc-rs on the resolver path, and rustls 0.23 refuses to auto-pick when either could win. Build (build.rs, Cargo.toml): tonic-build generates Rust types from proto/salience.proto at cargo-build time, using a vendored protoc binary (protoc-bin-vendored) so no system install is required. New runtime deps: tonic, prost, async-stream, tokio-stream, rustls-pemfile. Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
parent
0e459aae92
commit
08213f9093
15 changed files with 1689 additions and 440 deletions
|
|
@ -100,7 +100,7 @@ impl HttpClient {
|
|||
.map_err(|e| anyhow::anyhow!("invalid server name: {e}"))?;
|
||||
let connector = tokio_rustls::TlsConnector::from(self.tls.clone());
|
||||
let tls = connector.connect(server_name.to_owned(), tcp).await
|
||||
.context("TLS handshake")?;
|
||||
.map_err(|e| anyhow::anyhow!("TLS handshake to {host}: {e}"))?;
|
||||
TokioIo::new(Box::new(tls) as Box<dyn IoStream>)
|
||||
} else {
|
||||
TokioIo::new(Box::new(tcp) as Box<dyn IoStream>)
|
||||
|
|
@ -190,6 +190,7 @@ impl HttpClientBuilder {
|
|||
}
|
||||
|
||||
pub fn build(self) -> HttpClient {
|
||||
install_rustls_crypto_provider();
|
||||
let certs = rustls_native_certs::load_native_certs()
|
||||
.certs.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
|
|
@ -197,6 +198,13 @@ impl HttpClientBuilder {
|
|||
for cert in certs {
|
||||
root_store.add(cert).ok();
|
||||
}
|
||||
// Also trust any `.pem` files under `~/.consciousness/certs/` —
|
||||
// self-signed server certs for our own vllm hosts live there.
|
||||
// Drop a new `<host>.pem` in the dir to trust a new server; no
|
||||
// code change needed.
|
||||
for cert in load_user_certs() {
|
||||
root_store.add(cert).ok();
|
||||
}
|
||||
let tls = Arc::new(
|
||||
ClientConfig::builder()
|
||||
.with_root_certificates(root_store)
|
||||
|
|
@ -210,6 +218,65 @@ impl HttpClientBuilder {
|
|||
}
|
||||
}
|
||||
|
||||
/// Install rustls' default crypto provider exactly once per process.
|
||||
/// rustls 0.23 doesn't pick one automatically when multiple features
|
||||
/// could provide it (e.g. when tonic pulls in both ring and aws-lc-rs
|
||||
/// via transitive deps). Idempotent via OnceLock; safe to call from
|
||||
/// multiple callers.
|
||||
fn install_rustls_crypto_provider() {
|
||||
static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
|
||||
ONCE.get_or_init(|| {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
});
|
||||
}
|
||||
|
||||
/// Load every `.pem` file under `~/.consciousness/certs/` as a DER
|
||||
/// certificate and return them. Silent on missing dir, missing files,
|
||||
/// or parse errors — those are "no extra certs trusted" rather than
|
||||
/// hard failures, to keep startup robust.
|
||||
/// Load the concatenated PEM bytes of every `.pem` file under
|
||||
/// `~/.consciousness/certs/` — suitable for passing to a tonic
|
||||
/// `ClientTlsConfig::ca_certificate(Certificate::from_pem(...))` call
|
||||
/// so gRPC connections trust the same self-signed servers the HTTP
|
||||
/// path does.
|
||||
pub(crate) fn load_user_certs_pem_bytes() -> Vec<u8> {
|
||||
let mut out = Vec::new();
|
||||
let Some(home) = dirs::home_dir() else { return out };
|
||||
let dir = home.join(".consciousness").join("certs");
|
||||
let Ok(entries) = std::fs::read_dir(&dir) else { return out };
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("pem") {
|
||||
continue;
|
||||
}
|
||||
if let Ok(bytes) = std::fs::read(&path) {
|
||||
out.extend_from_slice(&bytes);
|
||||
if !bytes.ends_with(b"\n") {
|
||||
out.push(b'\n');
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn load_user_certs() -> Vec<rustls::pki_types::CertificateDer<'static>> {
|
||||
let mut out = Vec::new();
|
||||
let Some(home) = dirs::home_dir() else { return out };
|
||||
let dir = home.join(".consciousness").join("certs");
|
||||
let Ok(entries) = std::fs::read_dir(&dir) else { return out };
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("pem") {
|
||||
continue;
|
||||
}
|
||||
let Ok(bytes) = std::fs::read(&path) else { continue };
|
||||
for cert in rustls_pemfile::certs(&mut bytes.as_slice()).flatten() {
|
||||
out.push(cert);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Trait alias for streams that work with hyper's IO adapter.
|
||||
trait IoStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + 'static {}
|
||||
impl<T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + 'static> IoStream for T {}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue