use anyhow::{Context, Result};
use capnweb_core::CapId;
use reqwest::Client as HttpClient;
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, trace};
#[derive(Debug, Clone)]
pub struct ClientConfig {
pub url: String,
pub max_batch_size: usize,
pub timeout_ms: u64,
}
impl Default for ClientConfig {
fn default() -> Self {
Self {
url: "http://localhost:3000/rpc/batch".to_string(),
max_batch_size: 100,
timeout_ms: 30000,
}
}
}
pub struct Client {
config: ClientConfig,
http_client: HttpClient,
next_call_id: AtomicU64,
capabilities: Arc<RwLock<HashMap<CapId, Value>>>,
}
impl Client {
pub fn new(config: ClientConfig) -> Result<Self> {
let http_client = HttpClient::builder()
.timeout(std::time::Duration::from_millis(config.timeout_ms))
.build()
.context("Failed to build HTTP client")?;
Ok(Self {
config,
http_client,
next_call_id: AtomicU64::new(1), capabilities: Arc::new(RwLock::new(HashMap::new())),
})
}
pub fn new_with_url(url: &str) -> Result<Self> {
let config = ClientConfig {
url: url.to_string(),
..Default::default()
};
Self::new(config)
}
pub fn batch(&self) -> BatchBuilder<'_> {
BatchBuilder::new(self)
}
pub async fn call(&self, cap_id: CapId, method: &str, args: Vec<Value>) -> Result<Value> {
let import_id = self.next_call_id.fetch_add(1, Ordering::SeqCst);
let messages = vec![
json!(["push", ["call", cap_id.as_u64(), [method], args]]),
json!(["pull", import_id]),
];
let results = self.send_batch(messages).await?;
for result in results {
if let Some(arr) = result.as_array() {
if arr.len() >= 3
&& (arr[0] == "result" || arr[0] == "resolve")
&& arr[1] == import_id
{
return Ok(arr[2].clone());
} else if arr.len() >= 3 && arr[0] == "error" && arr[1] == import_id {
let error = arr[2]
.as_object()
.and_then(|o| o.get("message"))
.and_then(|m| m.as_str())
.unwrap_or("Unknown error");
return Err(anyhow::anyhow!("RPC error: {}", error));
}
}
}
Err(anyhow::anyhow!("No result received for call"))
}
async fn send_batch(&self, messages: Vec<Value>) -> Result<Vec<Value>> {
let body_parts: Result<Vec<String>> = messages
.iter()
.map(|m| serde_json::to_string(m).context("Failed to serialize message"))
.collect();
let body = body_parts?.join("\n");
debug!(
"Sending batch request to {}: {} messages",
self.config.url,
messages.len()
);
trace!("Request body:\n{}", body);
let response = self
.http_client
.post(&self.config.url)
.header("Content-Type", "text/plain")
.body(body)
.send()
.await
.context("Failed to send HTTP request")?;
let status = response.status();
let text = response
.text()
.await
.context("Failed to read response body")?;
if !status.is_success() {
return Err(anyhow::anyhow!("HTTP error {}: {}", status, text));
}
trace!("Response body:\n{}", text);
let mut results = Vec::new();
for line in text.lines() {
let line = line.trim();
if !line.is_empty() {
let value: Value = serde_json::from_str(line)
.with_context(|| format!("Failed to parse response line: {}", line))?;
results.push(value);
}
}
debug!("Received {} response messages", results.len());
Ok(results)
}
pub async fn register_capability(&self, id: CapId, cap: Value) {
let mut caps = self.capabilities.write().await;
caps.insert(id, cap);
}
pub async fn get_capability(&self, id: CapId) -> Option<Value> {
let caps = self.capabilities.read().await;
caps.get(&id).cloned()
}
pub async fn dispose_capability(&self, id: CapId) -> Result<()> {
let messages = vec![json!(["dispose", id.as_u64()])];
self.send_batch(messages).await?;
let mut caps = self.capabilities.write().await;
caps.remove(&id);
Ok(())
}
}
pub struct BatchBuilder<'a> {
client: &'a Client,
operations: Vec<BatchOperation>,
next_result_id: u64,
}
#[derive(Debug, Clone)]
pub struct BatchOperation {
pub id: u64,
pub message: Value,
pub is_pipeline: bool,
pub depends_on: Option<u64>,
}
#[derive(Debug, Clone)]
pub struct PendingResult {
pub id: u64,
pub path: Vec<String>,
}
impl<'a> BatchBuilder<'a> {
fn new(client: &'a Client) -> Self {
Self {
client,
operations: Vec::new(),
next_result_id: 1, }
}
pub fn call(&mut self, cap_id: CapId, method: &str, args: Vec<Value>) -> PendingResult {
let result_id = self.next_result_id;
self.next_result_id += 1;
let message = json!(["push", ["call", cap_id.as_u64(), [method], args]]);
self.operations.push(BatchOperation {
id: result_id,
message,
is_pipeline: false,
depends_on: None,
});
PendingResult {
id: result_id,
path: vec![],
}
}
pub fn pipeline(
&mut self,
base: &PendingResult,
path: Vec<&str>,
method: &str,
args: Vec<Value>,
) -> PendingResult {
let result_id = self.next_result_id;
self.next_result_id += 1;
let mut pipeline_args = Vec::new();
if !path.is_empty() {
let path_strings: Vec<Value> = path.iter().map(|s| json!(s)).collect();
pipeline_args.push(json!(["pipeline", base.id, path_strings]));
}
pipeline_args.extend(args);
let message = json!([
"push",
["pipeline", base.id, vec![json!(method)], pipeline_args]
]);
self.operations.push(BatchOperation {
id: result_id,
message,
is_pipeline: true,
depends_on: Some(base.id),
});
PendingResult {
id: result_id,
path: vec![],
}
}
pub fn reference(&self, result: &PendingResult, field: &str) -> PendingResult {
PendingResult {
id: result.id,
path: {
let mut path = result.path.clone();
path.push(field.to_string());
path
},
}
}
pub async fn execute(self) -> Result<BatchResults> {
if self.operations.is_empty() {
return Ok(BatchResults {
results: HashMap::new(),
});
}
let mut messages = Vec::new();
for op in &self.operations {
messages.push(op.message.clone());
}
for op in &self.operations {
messages.push(json!(["pull", op.id]));
}
let responses = self.client.send_batch(messages).await?;
let mut results = HashMap::new();
for response in responses {
if let Some(arr) = response.as_array() {
if arr.len() >= 2 {
let msg_type = arr[0].as_str().unwrap_or("");
match msg_type {
"result" | "resolve" => {
if let Some(id) = arr[1].as_u64() {
let value = arr.get(2).cloned().unwrap_or(json!(null));
results.insert(id, Ok(value));
}
}
"error" => {
if let Some(id) = arr[1].as_u64() {
let error_obj = arr.get(2).cloned().unwrap_or(json!({}));
let error_msg = error_obj
.get("message")
.and_then(|m| m.as_str())
.unwrap_or("Unknown error");
results
.insert(id, Err(anyhow::anyhow!("RPC error: {}", error_msg)));
}
}
_ => {}
}
}
}
}
Ok(BatchResults { results })
}
}
pub struct BatchResults {
results: HashMap<u64, Result<Value>>,
}
impl BatchResults {
pub fn get(&self, pending: &PendingResult) -> Result<Value> {
let value = self
.results
.get(&pending.id)
.ok_or_else(|| anyhow::anyhow!("No result for operation {}", pending.id))?
.as_ref()
.map_err(|e| anyhow::anyhow!("{}", e))?
.clone();
if pending.path.is_empty() {
Ok(value)
} else {
let mut current = value;
for segment in &pending.path {
current = current
.get(segment)
.ok_or_else(|| anyhow::anyhow!("Field '{}' not found in result", segment))?
.clone();
}
Ok(current)
}
}
pub fn contains(&self, pending: &PendingResult) -> bool {
self.results.contains_key(&pending.id)
}
pub fn all(&self) -> &HashMap<u64, Result<Value>> {
&self.results
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_client_creation() {
let config = ClientConfig::default();
let client = Client::new(config);
assert!(client.is_ok());
}
#[test]
fn test_batch_builder() {
let config = ClientConfig::default();
let client = Client::new(config).unwrap();
let mut batch = client.batch();
let result = batch.call(CapId::new(1), "test", vec![json!("arg")]);
assert_eq!(result.id, 1);
let pipelined = batch.pipeline(&result, vec!["field"], "method", vec![]);
assert_eq!(pipelined.id, 2);
}
}