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    /// Sovereign tier (D3): classify a prompt under a named routing profile (`Fast`|`Expert`|`Heavy`|
77    /// `Code`). Rust has no default args, so this is a sibling of [`route`] rather than an optional
78    /// param; `profile` is the snake_case cross-SDK contract field the edge `resolveProfile()` honors.
79    pub fn route_with_profile(&self, prompt: &str, profile: &str) -> Result<Value, Box<dyn Error>> {
80        self.post(&self.endpoint, with_profile(json!({ "prompt": prompt }), profile))
81    }
82
83    /// Classify and run on the edge if your plan allows it.
84    pub fn execute(&self, prompt: &str) -> Result<Value, Box<dyn Error>> {
85        self.post(&self.endpoint, json!({ "prompt": prompt, "execute": true }))
86    }
87
88    /// Execute under a named routing profile (Sovereign tier; see [`route_with_profile`]).
89    pub fn execute_with_profile(&self, prompt: &str, profile: &str) -> Result<Value, Box<dyn Error>> {
90        self.post(&self.endpoint, with_profile(json!({ "prompt": prompt, "execute": true }), profile))
91    }
92
93    /// Classify a pre-computed 768-d embedding (matmul-only: cheapest + fastest).
94    pub fn route_vector(&self, vector: &[f32]) -> Result<Value, Box<dyn Error>> {
95        self.post(&self.endpoint, json!({ "vector": vector }))
96    }
97
98    /// This license's savings ledger (decisions, saved_usd, saved_pct, ...). Requires a license.
99    pub fn savings(&self) -> Result<Value, Box<dyn Error>> {
100        self.get(&format!("{}/ledger/summary?license={}", self.agents, self.lic()?))
101    }
102
103    /// This license's agent-subscription status.
104    pub fn subscription(&self) -> Result<Value, Box<dyn Error>> {
105        self.get(&format!("{}/subscription/status?license={}", self.agents, self.lic()?))
106    }
107
108    /// Start/replace a programmatic subscription (plan: agent_starter|agent_pro|agent_scale).
109    pub fn subscribe(&self, plan: &str) -> Result<Value, Box<dyn Error>> {
110        self.post(&format!("{}/subscription/create", self.agents),
111                  json!({ "license": self.license, "plan": plan }))
112    }
113
114    /// 0.5.1 — Build a `PaymentHook` that signs each 402 challenge via a wallet provider.
115    ///
116    /// `provider`: `"cdp"` | `"privy"` | `"bridge"`. For custom wallets, build your own
117    /// closure and pass it to `with_payment_hook` directly.
118    /// `credentials`: provider-specific. CDP: `{api_key, api_secret, address}`. Privy: `{app_id,
119    /// app_secret, wallet_id}`. Bridge: `{api_key, source_wallet, destination?}`.
120    pub fn wallet_hook(provider: &str, credentials: HashMap<String, String>) -> Result<PaymentHook, Box<dyn Error>> {
121        let p = provider.to_string();
122        match p.as_str() {
123            "cdp" | "privy" | "bridge" => {
124                let header_name: &'static str = match p.as_str() {
125                    "cdp"    => "cdp-payment",
126                    "privy"  => "privy-payment",
127                    "bridge" => "bridge-payment",
128                    _ => unreachable!(),
129                };
130                let creds = credentials;
131                let proto = p.clone();
132                let hook: PaymentHook = Box::new(move |challenge: &Value| -> Result<HashMap<String, String>, Box<dyn Error>> {
133                    let payload = wallet_sign(&proto, &creds, challenge)?;
134                    let mut h = HashMap::new();
135                    h.insert(header_name.to_string(), payload);
136                    Ok(h)
137                });
138                Ok(hook)
139            }
140            other => Err(format!("dispatch::wallet_hook: unknown provider \"{}\"", other).into()),
141        }
142    }
143
144    fn lic(&self) -> Result<String, Box<dyn Error>> {
145        self.license.clone().ok_or_else(|| "dispatch: a license is required for savings()/subscription()".into())
146    }
147
148    fn auth(&self, mut req: ureq::Request) -> ureq::Request {
149        if let Some(l) = &self.license {
150            req = req.set("authorization", &format!("Bearer {}", l));
151        }
152        req
153    }
154
155    fn post(&self, url: &str, body: Value) -> Result<Value, Box<dyn Error>> {
156        let body_str = body.to_string();
157        let req = self.auth(ureq::post(url).set("content-type", "application/json"));
158        match req.send_string(&body_str) {
159            Ok(r) => Ok(serde_json::from_str(&r.into_string()?)?),
160            Err(ureq::Error::Status(402, r)) => self.retry_with_hook("POST", url, Some(&body_str), r),
161            Err(e) => Err(Box::new(e)),
162        }
163    }
164
165    fn get(&self, url: &str) -> Result<Value, Box<dyn Error>> {
166        let req = self.auth(ureq::get(url).set("content-type", "application/json"));
167        match req.call() {
168            Ok(r) => Ok(serde_json::from_str(&r.into_string()?)?),
169            Err(ureq::Error::Status(402, r)) => self.retry_with_hook("GET", url, None, r),
170            Err(e) => Err(Box::new(e)),
171        }
172    }
173
174    fn retry_with_hook(&self, method: &str, url: &str, body: Option<&str>, r: ureq::Response) -> Result<Value, Box<dyn Error>> {
175        let hook = self.payment_hook.as_ref()
176            .ok_or("dispatch: 402 payment required (x402) — pay and retry, or set a license / payment_hook")?;
177        let challenge: Value = serde_json::from_str(&r.into_string()?)?;
178        let headers = hook(&challenge)?;
179        let mut retry = if method == "POST" {
180            self.auth(ureq::post(url).set("content-type", "application/json"))
181        } else {
182            self.auth(ureq::get(url).set("content-type", "application/json"))
183        };
184        for (k, v) in &headers {
185            retry = retry.set(k, v);
186        }
187        let resp = if let Some(b) = body { retry.send_string(b)? } else { retry.call()? };
188        Ok(serde_json::from_str(&resp.into_string()?)?)
189    }
190}
191
192// Sovereign tier (D3): attach a named routing profile to a request body. `profile` is the snake_case
193// cross-SDK contract field; the edge resolveProfile() honors body.profile over per-license KV + defaults.
194// Empty profile => unchanged body. Mutates+returns the JSON object in place.
195fn with_profile(mut body: Value, profile: &str) -> Value {
196    if !profile.is_empty() {
197        if let Value::Object(ref mut m) = body {
198            m.insert("profile".into(), Value::String(profile.to_string()));
199        }
200    }
201    body
202}
203
204// Built-in provider sign — HTTP orchestration only; actual signing happens at the provider.
205// CDP-JWT signing is non-trivial in pure Rust (needs P-256 ECDSA + JWT lib); the built-in returns a
206// marker payload that the worker accepts via the wave-payments adapter (when WAVE_VERIFY_URL is set).
207// For full on-chain CDP signing, build your own closure with the official Coinbase Rust SDK.
208fn wallet_sign(provider: &str, creds: &HashMap<String, String>, challenge: &Value) -> Result<String, Box<dyn Error>> {
209    let accepts = challenge.get("accepts").and_then(|a| a.as_array()).cloned().unwrap_or_default();
210    let accept = accepts.iter().find(|a| a.get("protocol").and_then(|p| p.as_str()) == Some(provider))
211        .or_else(|| accepts.first())
212        .cloned()
213        .unwrap_or(Value::Null);
214
215    match provider {
216        "cdp" => {
217            // 0.6.2 — real CDP-JWT (ES256/P-256) signing via the `p256` crate (added to deps).
218            let api_key = creds.get("api_key").ok_or("dispatch::wallet_hook(cdp): api_key required")?;
219            let api_secret = creds.get("api_secret").ok_or("dispatch::wallet_hook(cdp): api_secret required (PEM EC private key)")?;
220            let jwt = sign_cdp_jwt(api_key, api_secret, &accept)?;
221            Ok(json!({
222                "provider": "cdp",
223                "jwt": jwt,
224                "address": creds.get("address"),
225                "network": creds.get("network").map(|s| s.as_str()).unwrap_or("base"),
226                "accept": accept
227            }).to_string())
228        },
229        "privy" => {
230            let app_id = creds.get("app_id").ok_or("dispatch::wallet_hook(privy): app_id required")?;
231            let app_secret = creds.get("app_secret").ok_or("dispatch::wallet_hook(privy): app_secret required")?;
232            let wallet_id = creds.get("wallet_id").ok_or("dispatch::wallet_hook(privy): wallet_id required")?;
233            use base64::Engine;
234            let basic = base64::engine::general_purpose::STANDARD.encode(format!("{}:{}", app_id, app_secret).as_bytes());
235            let body = json!({ "method": "personal_sign", "params": { "message": accept.to_string() }, "chain_type": "ethereum" }).to_string();
236            let url = format!("https://auth.privy.io/api/v1/wallets/{}/rpc", urlencoding::encode(wallet_id));
237            let resp = ureq::post(&url)
238                .set("content-type", "application/json")
239                .set("authorization", &format!("Basic {}", basic))
240                .set("privy-app-id", app_id)
241                .send_string(&body)?;
242            let j: Value = serde_json::from_str(&resp.into_string()?)?;
243            let sig = j.get("data").and_then(|d| d.get("signature")).or_else(|| j.get("signature")).cloned().unwrap_or(Value::Null);
244            Ok(json!({ "provider": "privy", "signature": sig, "accept": accept }).to_string())
245        }
246        "bridge" => {
247            let api_key = creds.get("api_key").ok_or("dispatch::wallet_hook(bridge): api_key required")?;
248            let body = json!({
249                "source": creds.get("source_wallet"),
250                "destination": creds.get("destination").cloned().unwrap_or_else(|| accept.get("payTo").cloned().unwrap_or(Value::Null).as_str().map(String::from).unwrap_or_default()),
251                "amount": accept.get("maxAmountRequired")
252            }).to_string();
253            let resp = ureq::post("https://api.bridge.xyz/v0/transfers")
254                .set("content-type", "application/json")
255                .set("api-key", api_key)
256                .send_string(&body)?;
257            let j: Value = serde_json::from_str(&resp.into_string()?)?;
258            Ok(json!({ "provider": "bridge", "id": j.get("id"), "accept": accept }).to_string())
259        }
260        other => Err(format!("dispatch::wallet_sign: unsupported provider \"{}\"", other).into()),
261    }
262}