Rust Client

The Rust client uses Tokio + async-tungstenite under the hood. Public types live in the pie_client crate.

Installation

[dependencies]
pie-client = "*"
tokio = { version = "1", features = ["full"] }
anyhow = "1"

Quick start

use pie_client::{Client, ParsedPrivateKey, ProcessEvent};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let client = Client::connect("ws://127.0.0.1:8080").await?;

    // Public-key auth (Ed25519/RSA/ECDSA OpenSSH or PEM).
    let key = ParsedPrivateKey::from_file("~/.ssh/id_ed25519")?;
    client.authenticate("alice", &Some(key)).await?;

    // input is a JSON string. Use serde_json::to_string(&MyInput { … })?.
    let mut process = client
        .launch_process(
            "text-completion@0.1.0".into(),
            r#"{"prompt":"Hello","max_tokens":64}"#.into(),
            true,
            None,
        )
        .await?;

    loop {
        match process.recv().await? {
            ProcessEvent::Stdout(s) => print!("{s}"),
            ProcessEvent::Return(v) => { println!("\n[done] {v}"); break; }
            ProcessEvent::Error(e)  => { eprintln!("\n[error] {e}"); break; }
            _ => {}
        }
    }

    client.close().await?;
    Ok(())
}

Client

Created with Client::connect(ws_host).await. Internally it spawns reader and writer tasks; they are torn down by client.close().await.

Authentication

impl Client {
    pub async fn authenticate(
        &self,
        username: &str,
        private_key: &Option<ParsedPrivateKey>,
    ) -> Result<()>;

    pub async fn auth_by_token(&self, token: &str) -> Result<()>;
}

authenticate runs a challenge-response handshake. auth_by_token is for backend ↔ engine internal connections. If the server has [auth].enabled = false, authenticate returns Ok(()) even with a None key.

Programs

pub async fn check_program(
    &self,
    inferlet: &str,                       // "name@major.minor.patch"
    wasm_path: Option<&Path>,
    manifest_path: Option<&Path>,
) -> Result<bool>;

pub async fn add_program(
    &self,
    wasm_path: &Path,
    manifest_path: &Path,
    force_overwrite: bool,
) -> Result<()>;

check_program validates name@version. When both paths are given, hashes are sent so the server can confirm the local copy matches what's installed. program_exists is a backwards-compatibility alias for check_program.

add_program uploads a local build in 256 KiB chunks.

Processes

pub async fn launch_process(
    &self,
    inferlet: String,
    input: String,                  // JSON
    capture_outputs: bool,
    token_budget: Option<usize>,
) -> Result<Process>;

pub async fn attach_process(&self, process_id: &str) -> Result<Process>;

pub async fn list_processes(&self) -> Result<Vec<String>>;

pub async fn terminate_process(&self, process_id: &str) -> Result<()>;

pub async fn ping(&self) -> Result<()>;

pub async fn query<T: ToString>(&self, subject: T, record: String) -> Result<String>;

pub async fn register_mcp_server(
    &self,
    name: &str,
    transport: &str,
    command: Option<&str>,
    args: Option<Vec<String>>,
    url: Option<&str>,
) -> Result<()>;

register_mcp_server exposes a local MCP server to inferlets in this session. The MCP request side of the bridge is currently a stub. See handle_server_message in client/rust/src/client.rs.

Process

pub struct Process { /* … */ }

impl Process {
    pub fn id(&self) -> &str;

    pub async fn signal<T: ToString>(&self, message: T) -> Result<()>;
    pub async fn transfer_file(&self, blob: &[u8]) -> Result<()>;
    pub async fn recv(&mut self) -> Result<ProcessEvent>;
    pub fn try_recv(&mut self) -> Result<Option<ProcessEvent>>;
}

Method	Description
id()	The process UUID.
signal(msg)	Fire-and-forget string signal to the running inferlet.
transfer_file(bytes)	Upload bytes to the process (chunked).
recv()	Await the next event.
try_recv()	Non-blocking variant: returns None when no event is queued.

ProcessEvent

pub enum ProcessEvent {
    Stdout(String),
    Stderr(String),
    Message(String),
    File(Vec<u8>),
    Return(String),     // JSON-encoded Output; process exits cleanly after this
    Error(String),      // process aborted
}

Match on the variant to dispatch:

match process.recv().await? {
    ProcessEvent::Stdout(s)  => print!("{s}"),
    ProcessEvent::Stderr(s)  => eprint!("{s}"),
    ProcessEvent::Message(m) => log::info!("inferlet: {m}"),
    ProcessEvent::File(b)    => save_blob(&b),
    ProcessEvent::Return(v)  => return Ok(serde_json::from_str(&v)?),
    ProcessEvent::Error(e)   => anyhow::bail!("process error: {e}"),
}

ParsedPrivateKey

Loads SSH/PEM private keys for authenticate. Supports OpenSSH (RSA / Ed25519 / ECDSA), PKCS#8 PEM, PKCS#1 PEM. ECDSA curves: P-256, P-384. RSA keys must be ≥ 2048 bits.

let key = ParsedPrivateKey::from_file("~/.ssh/id_ed25519")?;
// or
let key = ParsedPrivateKey::parse(&pem_string)?;

Detached / reattach

Process is just a handle. Drop the Client, reconnect, and reattach by id:

let pid = process.id().to_string();
client.close().await?;

// later …
let client = Client::connect(uri).await?;
client.authenticate("alice", &Some(key)).await?;
let mut process = client.attach_process(&pid).await?;
let event = process.recv().await?;

While disconnected, the server buffers events for you.

Utilities

pub fn hash_blob(blob: &[u8]) -> String;

Returns the hex blake3 hash. Useful for pre-computing program_hash when talking to the server directly.