Skip to main content

wave_dispatch/
lib.rs

1//! wave Dispatch — thin Rust client. Route each request to the cheapest capable model (local-first;
2//! escalate to your frontier only when needed). BYO keys + infra; the service returns a routing decision.
3use serde_json::{json, Value};
4use std::collections::HashMap;
5use std::error::Error;
6use std::time::{SystemTime, UNIX_EPOCH};
7
8/// 0.6.2 — CDP-JWT (ES256/P-256) signer. Uses the `p256` crate (added to dependencies). Pure Rust;
9/// no openssl FFI. Header: {alg:'ES256', kid:<api_key>, typ:'JWT', nonce:<rand-hex16>}. Payload:
10/// {sub:<api_key>, iss:'cdp', nbf:<now>, exp:<now+120>, uri:'POST dispatch.wave.online<resource>', claim:<accept>}.
11/// p256's ECDSA signer returns r||s as a Signature; we serialize to raw 32+32 (IEEE P-1363) — the JWS
12/// ES256 spec representation — then base64url-encode the three parts.
13pub fn sign_cdp_jwt(api_key: &str, pem_secret: &str, accept: &Value) -> Result<String, Box<dyn Error>> {
14    use base64::Engine;
15    use p256::ecdsa::{signature::Signer, Signature, SigningKey};
16    use p256::pkcs8::DecodePrivateKey;
17
18    let now = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0) as i64;
19    // 16 random bytes → 32 hex chars (uses rand from p256/sec1 transitive)
20    let mut nonce_bytes = [0u8; 16];
21    {
22        use rand_core::{OsRng, RngCore};
23        OsRng.fill_bytes(&mut nonce_bytes);
24    }
25    let nonce = nonce_bytes.iter().fold(String::with_capacity(32), |mut s, b| { use std::fmt::Write; let _ = write!(s, "{:02x}", b); s });
26    let resource = accept.get("resource").and_then(|v| v.as_str()).unwrap_or("/");
27    let uri = format!("POST dispatch.wave.online{}", resource);
28    let header = json!({ "alg": "ES256", "kid": api_key, "typ": "JWT", "nonce": nonce });
29    let payload = json!({ "sub": api_key, "iss": "cdp", "nbf": now, "exp": now + 120, "uri": uri, "claim": accept });
30    let b64 = |s: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(s);
31    let to_sign = format!("{}.{}", b64(header.to_string().as_bytes()), b64(payload.to_string().as_bytes()));
32
33    let key = SigningKey::from_pkcs8_pem(pem_secret)
34        .map_err(|e| format!("dispatch::sign_cdp_jwt: PEM parse failed ({}); api_secret must be a PKCS8 EC P-256 private key", e))?;
35    let sig: Signature = key.sign(to_sign.as_bytes());
36    Ok(format!("{}.{}", to_sign, b64(sig.to_bytes().as_slice())))
37}
38
39/// 0.5.1 — payment hook: called once with the 402 challenge body, returns headers to retry the
40/// request with (e.g. `{"x-payment": "..."}` for x402, `{"tempo-payment": "..."}` for tempo). Pair with
41/// `Dispatch::wallet_hook(provider, credentials)` for the built-in CDP / Privy / Bridge factories, or
42/// build a closure yourself for any custom wallet stack.
43pub type PaymentHook =
44    Box<dyn Fn(&Value) -> Result<HashMap<String, String>, Box<dyn Error>> + Send + Sync>;
45
46pub struct Dispatch {
47    license: Option<String>,
48    endpoint: String,
49    agents: String,
50    payment_hook: Option<PaymentHook>,
51}
52
53impl Dispatch {
54    /// `license`: your `wv_...` key, or `None` to read `WAVE_LICENSE` (omit for x402 pay-per-use).
55    pub fn new(license: Option<String>) -> Self {
56        Dispatch {
57            license: license.or_else(|| std::env::var("WAVE_LICENSE").ok()),
58            endpoint: std::env::var("DISPATCH_ENDPOINT").unwrap_or_else(|_| "https://dispatch.wave.online".into()),
59            agents: std::env::var("WAVE_AGENTS_ENDPOINT").unwrap_or_else(|_| "https://dispatch-agents.wave.online".into()),
60            payment_hook: None,
61        }
62    }
63
64    /// 0.5.1 — attach a payment hook so 402 challenges are handled inside the client (signs + retries
65    /// in one .route() call). See `Dispatch::wallet_hook` for the built-in factory.
66    pub fn with_payment_hook(mut self, hook: PaymentHook) -> Self {
67        self.payment_hook = Some(hook);
68        self
69    }
70
71    /// Classify a prompt (no execution): `{route, probability, margin, forward}`.
72    pub fn route(&self, prompt: &str) -> Result<Value, Box<dyn Error>> {
73        self.post(&self.endpoint, json!({ "prompt": prompt }))
74    }
75
76    /// Classify and run on the edge if your plan allows it.
77    pub fn execute(&self, prompt: &str) -> Result<Value, Box<dyn Error>> {
78        self.post(&self.endpoint, json!({ "prompt": prompt, "execute": true }))
79    }
80
81    /// Classify a pre-computed 768-d embedding (matmul-only: cheapest + fastest).
82    pub fn route_vector(&self, vector: &[f32]) -> Result<Value, Box<dyn Error>> {
83        self.post(&self.endpoint, json!({ "vector": vector }))
84    }
85
86    /// This license's savings ledger (decisions, saved_usd, saved_pct, ...). Requires a license.
87    pub fn savings(&self) -> Result<Value, Box<dyn Error>> {
88        self.get(&format!("{}/ledger/summary?license={}", self.agents, self.lic()?))
89    }
90
91    /// This license's agent-subscription status.
92    pub fn subscription(&self) -> Result<Value, Box<dyn Error>> {
93        self.get(&format!("{}/subscription/status?license={}", self.agents, self.lic()?))
94    }
95
96    /// Start/replace a programmatic subscription (plan: agent_starter|agent_pro|agent_scale).
97    pub fn subscribe(&self, plan: &str) -> Result<Value, Box<dyn Error>> {
98        self.post(&format!("{}/subscription/create", self.agents),
99                  json!({ "license": self.license, "plan": plan }))
100    }
101
102    /// 0.5.1 — Build a `PaymentHook` that signs each 402 challenge via a wallet provider.
103    ///
104    /// `provider`: `"cdp"` | `"privy"` | `"bridge"`. For custom wallets, build your own
105    /// closure and pass it to `with_payment_hook` directly.
106    /// `credentials`: provider-specific. CDP: `{api_key, api_secret, address}`. Privy: `{app_id,
107    /// app_secret, wallet_id}`. Bridge: `{api_key, source_wallet, destination?}`.
108    pub fn wallet_hook(provider: &str, credentials: HashMap<String, String>) -> Result<PaymentHook, Box<dyn Error>> {
109        let p = provider.to_string();
110        match p.as_str() {
111            "cdp" | "privy" | "bridge" => {
112                let header_name: &'static str = match p.as_str() {
113                    "cdp"    => "cdp-payment",
114                    "privy"  => "privy-payment",
115                    "bridge" => "bridge-payment",
116                    _ => unreachable!(),
117                };
118                let creds = credentials;
119                let proto = p.clone();
120                let hook: PaymentHook = Box::new(move |challenge: &Value| -> Result<HashMap<String, String>, Box<dyn Error>> {
121                    let payload = wallet_sign(&proto, &creds, challenge)?;
122                    let mut h = HashMap::new();
123                    h.insert(header_name.to_string(), payload);
124                    Ok(h)
125                });
126                Ok(hook)
127            }
128            other => Err(format!("dispatch::wallet_hook: unknown provider \"{}\"", other).into()),
129        }
130    }
131
132    fn lic(&self) -> Result<String, Box<dyn Error>> {
133        self.license.clone().ok_or_else(|| "dispatch: a license is required for savings()/subscription()".into())
134    }
135
136    fn auth(&self, mut req: ureq::Request) -> ureq::Request {
137        if let Some(l) = &self.license {
138            req = req.set("authorization", &format!("Bearer {}", l));
139        }
140        req
141    }
142
143    fn post(&self, url: &str, body: Value) -> Result<Value, Box<dyn Error>> {
144        let body_str = body.to_string();
145        let req = self.auth(ureq::post(url).set("content-type", "application/json"));
146        match req.send_string(&body_str) {
147            Ok(r) => Ok(serde_json::from_str(&r.into_string()?)?),
148            Err(ureq::Error::Status(402, r)) => self.retry_with_hook("POST", url, Some(&body_str), r),
149            Err(e) => Err(Box::new(e)),
150        }
151    }
152
153    fn get(&self, url: &str) -> Result<Value, Box<dyn Error>> {
154        let req = self.auth(ureq::get(url).set("content-type", "application/json"));
155        match req.call() {
156            Ok(r) => Ok(serde_json::from_str(&r.into_string()?)?),
157            Err(ureq::Error::Status(402, r)) => self.retry_with_hook("GET", url, None, r),
158            Err(e) => Err(Box::new(e)),
159        }
160    }
161
162    fn retry_with_hook(&self, method: &str, url: &str, body: Option<&str>, r: ureq::Response) -> Result<Value, Box<dyn Error>> {
163        let hook = self.payment_hook.as_ref()
164            .ok_or("dispatch: 402 payment required (x402) — pay and retry, or set a license / payment_hook")?;
165        let challenge: Value = serde_json::from_str(&r.into_string()?)?;
166        let headers = hook(&challenge)?;
167        let mut retry = if method == "POST" {
168            self.auth(ureq::post(url).set("content-type", "application/json"))
169        } else {
170            self.auth(ureq::get(url).set("content-type", "application/json"))
171        };
172        for (k, v) in &headers {
173            retry = retry.set(k, v);
174        }
175        let resp = if let Some(b) = body { retry.send_string(b)? } else { retry.call()? };
176        Ok(serde_json::from_str(&resp.into_string()?)?)
177    }
178}
179
180// Built-in provider sign — HTTP orchestration only; actual signing happens at the provider.
181// CDP-JWT signing is non-trivial in pure Rust (needs P-256 ECDSA + JWT lib); the built-in returns a
182// marker payload that the worker accepts via the wave-payments adapter (when WAVE_VERIFY_URL is set).
183// For full on-chain CDP signing, build your own closure with the official Coinbase Rust SDK.
184fn wallet_sign(provider: &str, creds: &HashMap<String, String>, challenge: &Value) -> Result<String, Box<dyn Error>> {
185    let accepts = challenge.get("accepts").and_then(|a| a.as_array()).cloned().unwrap_or_default();
186    let accept = accepts.iter().find(|a| a.get("protocol").and_then(|p| p.as_str()) == Some(provider))
187        .or_else(|| accepts.first())
188        .cloned()
189        .unwrap_or(Value::Null);
190
191    match provider {
192        "cdp" => {
193            // 0.6.2 — real CDP-JWT (ES256/P-256) signing via the `p256` crate (added to deps).
194            let api_key = creds.get("api_key").ok_or("dispatch::wallet_hook(cdp): api_key required")?;
195            let api_secret = creds.get("api_secret").ok_or("dispatch::wallet_hook(cdp): api_secret required (PEM EC private key)")?;
196            let jwt = sign_cdp_jwt(api_key, api_secret, &accept)?;
197            Ok(json!({
198                "provider": "cdp",
199                "jwt": jwt,
200                "address": creds.get("address"),
201                "network": creds.get("network").map(|s| s.as_str()).unwrap_or("base"),
202                "accept": accept
203            }).to_string())
204        },
205        "privy" => {
206            let app_id = creds.get("app_id").ok_or("dispatch::wallet_hook(privy): app_id required")?;
207            let app_secret = creds.get("app_secret").ok_or("dispatch::wallet_hook(privy): app_secret required")?;
208            let wallet_id = creds.get("wallet_id").ok_or("dispatch::wallet_hook(privy): wallet_id required")?;
209            use base64::Engine;
210            let basic = base64::engine::general_purpose::STANDARD.encode(format!("{}:{}", app_id, app_secret).as_bytes());
211            let body = json!({ "method": "personal_sign", "params": { "message": accept.to_string() }, "chain_type": "ethereum" }).to_string();
212            let url = format!("https://auth.privy.io/api/v1/wallets/{}/rpc", urlencoding::encode(wallet_id));
213            let resp = ureq::post(&url)
214                .set("content-type", "application/json")
215                .set("authorization", &format!("Basic {}", basic))
216                .set("privy-app-id", app_id)
217                .send_string(&body)?;
218            let j: Value = serde_json::from_str(&resp.into_string()?)?;
219            let sig = j.get("data").and_then(|d| d.get("signature")).or_else(|| j.get("signature")).cloned().unwrap_or(Value::Null);
220            Ok(json!({ "provider": "privy", "signature": sig, "accept": accept }).to_string())
221        }
222        "bridge" => {
223            let api_key = creds.get("api_key").ok_or("dispatch::wallet_hook(bridge): api_key required")?;
224            let body = json!({
225                "source": creds.get("source_wallet"),
226                "destination": creds.get("destination").cloned().unwrap_or_else(|| accept.get("payTo").cloned().unwrap_or(Value::Null).as_str().map(String::from).unwrap_or_default()),
227                "amount": accept.get("maxAmountRequired")
228            }).to_string();
229            let resp = ureq::post("https://api.bridge.xyz/v0/transfers")
230                .set("content-type", "application/json")
231                .set("api-key", api_key)
232                .send_string(&body)?;
233            let j: Value = serde_json::from_str(&resp.into_string()?)?;
234            Ok(json!({ "provider": "bridge", "id": j.get("id"), "accept": accept }).to_string())
235        }
236        other => Err(format!("dispatch::wallet_sign: unsupported provider \"{}\"", other).into()),
237    }
238}