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;
6
7/// 0.5.1 — payment hook: called once with the 402 challenge body, returns headers to retry the
8/// request with (e.g. `{"x-payment": "..."}` for x402, `{"tempo-payment": "..."}` for tempo). Pair with
9/// `Dispatch::wallet_hook(provider, credentials)` for the built-in CDP / Privy / Bridge factories, or
10/// build a closure yourself for any custom wallet stack.
11pub type PaymentHook =
12    Box<dyn Fn(&Value) -> Result<HashMap<String, String>, Box<dyn Error>> + Send + Sync>;
13
14pub struct Dispatch {
15    license: Option<String>,
16    endpoint: String,
17    agents: String,
18    payment_hook: Option<PaymentHook>,
19}
20
21impl Dispatch {
22    /// `license`: your `wv_...` key, or `None` to read `WAVE_LICENSE` (omit for x402 pay-per-use).
23    pub fn new(license: Option<String>) -> Self {
24        Dispatch {
25            license: license.or_else(|| std::env::var("WAVE_LICENSE").ok()),
26            endpoint: std::env::var("DISPATCH_ENDPOINT").unwrap_or_else(|_| "https://dispatch.wave.online".into()),
27            agents: std::env::var("WAVE_AGENTS_ENDPOINT").unwrap_or_else(|_| "https://dispatch-agents.wave.online".into()),
28            payment_hook: None,
29        }
30    }
31
32    /// 0.5.1 — attach a payment hook so 402 challenges are handled inside the client (signs + retries
33    /// in one .route() call). See `Dispatch::wallet_hook` for the built-in factory.
34    pub fn with_payment_hook(mut self, hook: PaymentHook) -> Self {
35        self.payment_hook = Some(hook);
36        self
37    }
38
39    /// Classify a prompt (no execution): `{route, probability, margin, forward}`.
40    pub fn route(&self, prompt: &str) -> Result<Value, Box<dyn Error>> {
41        self.post(&self.endpoint, json!({ "prompt": prompt }))
42    }
43
44    /// Classify and run on the edge if your plan allows it.
45    pub fn execute(&self, prompt: &str) -> Result<Value, Box<dyn Error>> {
46        self.post(&self.endpoint, json!({ "prompt": prompt, "execute": true }))
47    }
48
49    /// Classify a pre-computed 768-d embedding (matmul-only: cheapest + fastest).
50    pub fn route_vector(&self, vector: &[f32]) -> Result<Value, Box<dyn Error>> {
51        self.post(&self.endpoint, json!({ "vector": vector }))
52    }
53
54    /// This license's savings ledger (decisions, saved_usd, saved_pct, ...). Requires a license.
55    pub fn savings(&self) -> Result<Value, Box<dyn Error>> {
56        self.get(&format!("{}/ledger/summary?license={}", self.agents, self.lic()?))
57    }
58
59    /// This license's agent-subscription status.
60    pub fn subscription(&self) -> Result<Value, Box<dyn Error>> {
61        self.get(&format!("{}/subscription/status?license={}", self.agents, self.lic()?))
62    }
63
64    /// Start/replace a programmatic subscription (plan: agent_starter|agent_pro|agent_scale).
65    pub fn subscribe(&self, plan: &str) -> Result<Value, Box<dyn Error>> {
66        self.post(&format!("{}/subscription/create", self.agents),
67                  json!({ "license": self.license, "plan": plan }))
68    }
69
70    /// 0.5.1 — Build a `PaymentHook` that signs each 402 challenge via a wallet provider.
71    ///
72    /// `provider`: `"cdp"` | `"privy"` | `"bridge"`. For custom wallets, build your own
73    /// closure and pass it to `with_payment_hook` directly.
74    /// `credentials`: provider-specific. CDP: `{api_key, api_secret, address}`. Privy: `{app_id,
75    /// app_secret, wallet_id}`. Bridge: `{api_key, source_wallet, destination?}`.
76    pub fn wallet_hook(provider: &str, credentials: HashMap<String, String>) -> Result<PaymentHook, Box<dyn Error>> {
77        let p = provider.to_string();
78        match p.as_str() {
79            "cdp" | "privy" | "bridge" => {
80                let header_name: &'static str = match p.as_str() {
81                    "cdp"    => "cdp-payment",
82                    "privy"  => "privy-payment",
83                    "bridge" => "bridge-payment",
84                    _ => unreachable!(),
85                };
86                let creds = credentials;
87                let proto = p.clone();
88                let hook: PaymentHook = Box::new(move |challenge: &Value| -> Result<HashMap<String, String>, Box<dyn Error>> {
89                    let payload = wallet_sign(&proto, &creds, challenge)?;
90                    let mut h = HashMap::new();
91                    h.insert(header_name.to_string(), payload);
92                    Ok(h)
93                });
94                Ok(hook)
95            }
96            other => Err(format!("dispatch::wallet_hook: unknown provider \"{}\"", other).into()),
97        }
98    }
99
100    fn lic(&self) -> Result<String, Box<dyn Error>> {
101        self.license.clone().ok_or_else(|| "dispatch: a license is required for savings()/subscription()".into())
102    }
103
104    fn auth(&self, mut req: ureq::Request) -> ureq::Request {
105        if let Some(l) = &self.license {
106            req = req.set("authorization", &format!("Bearer {}", l));
107        }
108        req
109    }
110
111    fn post(&self, url: &str, body: Value) -> Result<Value, Box<dyn Error>> {
112        let body_str = body.to_string();
113        let req = self.auth(ureq::post(url).set("content-type", "application/json"));
114        match req.send_string(&body_str) {
115            Ok(r) => Ok(serde_json::from_str(&r.into_string()?)?),
116            Err(ureq::Error::Status(402, r)) => self.retry_with_hook("POST", url, Some(&body_str), r),
117            Err(e) => Err(Box::new(e)),
118        }
119    }
120
121    fn get(&self, url: &str) -> Result<Value, Box<dyn Error>> {
122        let req = self.auth(ureq::get(url).set("content-type", "application/json"));
123        match req.call() {
124            Ok(r) => Ok(serde_json::from_str(&r.into_string()?)?),
125            Err(ureq::Error::Status(402, r)) => self.retry_with_hook("GET", url, None, r),
126            Err(e) => Err(Box::new(e)),
127        }
128    }
129
130    fn retry_with_hook(&self, method: &str, url: &str, body: Option<&str>, r: ureq::Response) -> Result<Value, Box<dyn Error>> {
131        let hook = self.payment_hook.as_ref()
132            .ok_or("dispatch: 402 payment required (x402) — pay and retry, or set a license / payment_hook")?;
133        let challenge: Value = serde_json::from_str(&r.into_string()?)?;
134        let headers = hook(&challenge)?;
135        let mut retry = if method == "POST" {
136            self.auth(ureq::post(url).set("content-type", "application/json"))
137        } else {
138            self.auth(ureq::get(url).set("content-type", "application/json"))
139        };
140        for (k, v) in &headers {
141            retry = retry.set(k, v);
142        }
143        let resp = if let Some(b) = body { retry.send_string(b)? } else { retry.call()? };
144        Ok(serde_json::from_str(&resp.into_string()?)?)
145    }
146}
147
148// Built-in provider sign — HTTP orchestration only; actual signing happens at the provider.
149// CDP-JWT signing is non-trivial in pure Rust (needs P-256 ECDSA + JWT lib); the built-in returns a
150// marker payload that the worker accepts via the wave-payments adapter (when WAVE_VERIFY_URL is set).
151// For full on-chain CDP signing, build your own closure with the official Coinbase Rust SDK.
152fn wallet_sign(provider: &str, creds: &HashMap<String, String>, challenge: &Value) -> Result<String, Box<dyn Error>> {
153    let accepts = challenge.get("accepts").and_then(|a| a.as_array()).cloned().unwrap_or_default();
154    let accept = accepts.iter().find(|a| a.get("protocol").and_then(|p| p.as_str()) == Some(provider))
155        .or_else(|| accepts.first())
156        .cloned()
157        .unwrap_or(Value::Null);
158
159    match provider {
160        "cdp" => Ok(json!({
161            "provider": "cdp",
162            "address": creds.get("address"),
163            "accept": accept,
164            "hint": "use the official Coinbase Rust SDK for CDP-JWT signing in production"
165        }).to_string()),
166        "privy" => {
167            let app_id = creds.get("app_id").ok_or("dispatch::wallet_hook(privy): app_id required")?;
168            let app_secret = creds.get("app_secret").ok_or("dispatch::wallet_hook(privy): app_secret required")?;
169            let wallet_id = creds.get("wallet_id").ok_or("dispatch::wallet_hook(privy): wallet_id required")?;
170            use base64::Engine;
171            let basic = base64::engine::general_purpose::STANDARD.encode(format!("{}:{}", app_id, app_secret).as_bytes());
172            let body = json!({ "method": "personal_sign", "params": { "message": accept.to_string() }, "chain_type": "ethereum" }).to_string();
173            let url = format!("https://auth.privy.io/api/v1/wallets/{}/rpc", urlencoding::encode(wallet_id));
174            let resp = ureq::post(&url)
175                .set("content-type", "application/json")
176                .set("authorization", &format!("Basic {}", basic))
177                .set("privy-app-id", app_id)
178                .send_string(&body)?;
179            let j: Value = serde_json::from_str(&resp.into_string()?)?;
180            let sig = j.get("data").and_then(|d| d.get("signature")).or_else(|| j.get("signature")).cloned().unwrap_or(Value::Null);
181            Ok(json!({ "provider": "privy", "signature": sig, "accept": accept }).to_string())
182        }
183        "bridge" => {
184            let api_key = creds.get("api_key").ok_or("dispatch::wallet_hook(bridge): api_key required")?;
185            let body = json!({
186                "source": creds.get("source_wallet"),
187                "destination": creds.get("destination").cloned().unwrap_or_else(|| accept.get("payTo").cloned().unwrap_or(Value::Null).as_str().map(String::from).unwrap_or_default()),
188                "amount": accept.get("maxAmountRequired")
189            }).to_string();
190            let resp = ureq::post("https://api.bridge.xyz/v0/transfers")
191                .set("content-type", "application/json")
192                .set("api-key", api_key)
193                .send_string(&body)?;
194            let j: Value = serde_json::from_str(&resp.into_string()?)?;
195            Ok(json!({ "provider": "bridge", "id": j.get("id"), "accept": accept }).to_string())
196        }
197        other => Err(format!("dispatch::wallet_sign: unsupported provider \"{}\"", other).into()),
198    }
199}