Rust Client

A Rust client for interacting with the Pie server.

Installation

Add to your Cargo.toml:

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

Quick Start

use pie_client::{Client, InstanceEvent};
use anyhow::Result;

#[tokio::main]
async fn main() -> Result<()> {
    // Connect to server
    let client = Client::connect("ws://127.0.0.1:8080").await?;
    
    // Authenticate
    client.authenticate("username", &None).await?;
    
    // Launch an inferlet from registry
    let mut instance = client.launch_instance_from_registry(
        "text-completion".to_string(),
        vec!["--prompt".to_string(), "Hello, world!".to_string()],
        false,  // not detached
    ).await?;
    
    // Receive output
    loop {
        match instance.recv().await? {
            InstanceEvent::Stdout(text) => print!("{}", text),
            InstanceEvent::Completed(_) => break,
            InstanceEvent::Exception(msg) => {
                eprintln!("Error: {}", msg);
                break;
            }
            _ => {}
        }
    }
    
    client.close()?;
    Ok(())
}

Client

The main client struct for connecting to a Pie server.

Construction

let client = Client::connect("ws://127.0.0.1:8080").await?;

Methods

Method	Returns	Description
connect(ws_host)	Result<Client>	Connect to server
close(self)	Result<()>	Close connection
authenticate(username, key)	Result<()>	Authenticate with server
internal_authenticate(token)	Result<()>	Internal token auth
upload_program(bytes)	Result<()>	Upload WASM program
program_exists(hash)	Result<bool>	Check if program exists
launch_instance(hash, args, detached)	Result<Instance>	Launch by hash
launch_instance_from_registry(name, args, detached)	Result<Instance>	Launch from registry
attach_instance(id)	Result<Instance>	Attach to existing
list_instances()	Result<Vec<InstanceInfo>>	List instances
terminate_instance(id)	Result<()>	Terminate instance
ping()	Result<()>	Check connectivity

Authentication

use pie_client::crypto::ParsedPrivateKey;

// With SSH key
let key = ParsedPrivateKey::from_file("~/.ssh/id_ed25519")?;
client.authenticate("username", &Some(key)).await?;

// Without key (when server auth is disabled)
client.authenticate("username", &None).await?;

Instance

Represents a running program instance.

Methods

Method	Returns	Description
id()	InstanceId	Get instance ID
send(message)	Result<()>	Send string message
upload_blob(bytes)	Result<()>	Upload binary data
recv()	Result<InstanceEvent>	Receive next event

Example

let mut instance = client.launch_instance_from_registry(
    "chat".to_string(),
    vec![],
    false,
).await?;

// Send a message
instance.send("What is 2 + 2?")?;

// Receive response
match instance.recv().await? {
    InstanceEvent::Message(text) => println!("{}", text),
    InstanceEvent::Stdout(text) => print!("{}", text),
    _ => {}
}

InstanceEvent

Events received from instances:

pub enum InstanceEvent {
    /// Text message from instance
    Message(String),
    /// Instance completed successfully
    Completed(String),
    /// Instance was aborted
    Aborted(String),
    /// Instance raised an exception
    Exception(String),
    /// Server-side error
    ServerError(String),
    /// Resource limit reached
    OutOfResources(String),
    /// Binary data
    Blob(Vec<u8>),
    /// Streaming stdout
    Stdout(String),
    /// Streaming stderr
    Stderr(String),
}

Pattern Matching

loop {
    match instance.recv().await? {
        InstanceEvent::Stdout(text) => {
            print!("{}", text);
            std::io::stdout().flush()?;
        }
        InstanceEvent::Stderr(text) => {
            eprint!("{}", text);
        }
        InstanceEvent::Completed(msg) => {
            println!("\nCompleted: {}", msg);
            break;
        }
        InstanceEvent::Exception(err) => {
            return Err(anyhow!("Inferlet error: {}", err));
        }
        InstanceEvent::ServerError(err) => {
            return Err(anyhow!("Server error: {}", err));
        }
        _ => {}
    }
}

Upload Custom Inferlet

use pie_client::hash_blob;

// Read and upload
let wasm_bytes = std::fs::read("my_inferlet.wasm")?;
client.upload_program(&wasm_bytes).await?;

// Get hash for launching
let program_hash = hash_blob(&wasm_bytes);

// Launch
let mut instance = client.launch_instance(
    program_hash,
    vec!["--arg".to_string(), "value".to_string()],
    false,
).await?;

Detached Instances

// Launch detached
let instance = client.launch_instance_from_registry(
    "long-task".to_string(),
    vec![],
    true,  // detached
).await?;
let instance_id = instance.id();

// Later, reattach
let mut attached = client.attach_instance(&instance_id.to_string()).await?;
let event = attached.recv().await?;

Error Handling

use anyhow::{Result, Context};

async fn run() -> Result<()> {
    let client = Client::connect("ws://127.0.0.1:8080")
        .await
        .context("Failed to connect to Pie server")?;
    
    client.authenticate("username", &None)
        .await
        .context("Authentication failed")?;
    
    // ...
    Ok(())
}

Thread Safety

The Client is thread-safe and can be shared across tasks:

use std::sync::Arc;

let client = Arc::new(Client::connect("ws://127.0.0.1:8080").await?);

let c1 = client.clone();
let h1 = tokio::spawn(async move {
    // Use c1
});

let c2 = client.clone();
let h2 = tokio::spawn(async move {
    // Use c2
});

h1.await?;
h2.await?;