AgentDB with Rust

let base_url =
    std::env::var("AGENTDB_URL").unwrap_or_else(|_| "https://app.agentdb.link".into());
let api_key = std::env::var("AGENTDB_API_KEY"); // a server-side secret: never in browser code
let workspace = std::env::var("AGENTDB_WORKSPACE"); // ws_...
let (Ok(api_key), Ok(workspace)) = (api_key, workspace) else {
    return Err("Set AGENTDB_API_KEY and AGENTDB_WORKSPACE".into());
};
let api = Api {
    url: format!(
        "{}/v1/workspaces/{}",
        base_url.trim_end_matches('/'),
        workspace
    ),
    key: api_key,
    http: Client::new(),
};

/// An AgentDB error: HTTP status plus the error envelope {code, message, retryable, details}.
#[derive(Debug)]
struct AgentDBError {
    status: u16,
    code: String,
    message: String,
    retryable: bool,
    details: Value,
}

impl fmt::Display for AgentDBError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} {}: {}", self.status, self.code, self.message)
    }
}

impl Error for AgentDBError {}

/// The optional parts of a call.
#[derive(Default)]
struct Request {
    json: Option<Value>,             // sent as application/json
    body: Option<Vec<u8>>,           // sent as application/octet-stream
    idempotency_key: Option<String>, // makes a write safe to retry
}

impl Request {
    fn json(value: Value) -> Self {
        Request {
            json: Some(value),
            ..Default::default()
        }
    }

    fn idempotency_key(mut self, key: impl Into<String>) -> Self {
        self.idempotency_key = Some(key.into());
        self
    }
}

/// The workspace API: its URL, the API key and one HTTP client.
struct Api {
    url: String,
    key: String,
    http: Client,
}

impl Api {
    /// One API call. Sends the key in the Authorization header. Retries 429 and 503 with Retry-After
    /// (or exponential backoff), and network errors when the call is safe to repeat (GET, or a write
    /// with an Idempotency-Key). Returns the raw response.
    fn send(&self, method: Method, path: &str, req: &Request) -> Result<Response, Box<dyn Error>> {
        let safe_to_repeat = method == Method::GET || req.idempotency_key.is_some();
        for attempt in 0..=4u32 {
            let backoff = Duration::from_millis(500 << attempt);
            let mut builder = self
                .http
                .request(method.clone(), format!("{}{}", self.url, path))
                .bearer_auth(&self.key);
            if let Some(json) = &req.json {
                builder = builder.json(json);
            }
            if let Some(body) = &req.body {
                builder = builder
                    .header(CONTENT_TYPE, "application/octet-stream")
                    .body(body.clone());
            }
            if let Some(key) = &req.idempotency_key {
                builder = builder.header("Idempotency-Key", key);
            }
            let res = match builder.send() {
                Ok(res) => res,
                Err(_) if safe_to_repeat && attempt < 4 => {
                    sleep(backoff);
                    continue;
                }
                Err(err) => return Err(err.into()),
            };
            let status = res.status().as_u16();
            if (status == 429 || status == 503) && attempt < 4 {
                let retry_after = res
                    .headers()
                    .get(RETRY_AFTER)
                    .and_then(|v| v.to_str().ok()?.parse::<u64>().ok());
                let jitter = Duration::from_millis(u64::from(
                    SystemTime::now()
                        .duration_since(UNIX_EPOCH)?
                        .subsec_millis()
                        % 250,
                ));
                sleep(match retry_after {
                    Some(seconds) if seconds > 0 => Duration::from_secs(seconds),
                    _ => backoff + jitter,
                });
                continue;
            }
            if !res.status().is_success() {
                let envelope: Value = res.json().unwrap_or(Value::Null);
                let error = &envelope["error"];
                return Err(Box::new(AgentDBError {
                    status,
                    code: error["code"].as_str().unwrap_or("HTTP_ERROR").to_string(),
                    message: error["message"].as_str().unwrap_or("").to_string(),
                    retryable: error["retryable"] == true,
                    details: error["details"].clone(),
                }));
            }
            return Ok(res);
        }
        unreachable!("the last attempt always returns")
    }

    /// One API call that returns the "data" part of the answer.
    fn call(&self, method: Method, path: &str, req: Request) -> Result<Value, Box<dyn Error>> {
        let mut answer: Value = self.send(method, path, &req)?.json()?;
        Ok(answer["data"].take())
    }
}

let schema = api.call(Method::GET, "/object-types", Request::default())?;
let names: Vec<&str> = schema["objects"]
    .as_array()
    .into_iter()
    .flatten()
    .filter_map(|o| o["name"].as_str())
    .collect();
println!("Objects: {}", names.join(", "));

let ada = api.call(
    Method::POST,
    "/records",
    Request::json(json!({ "object": "contact", "data": { "name": "Ada Lovelace", "email": format!("ada-{run}@example.com"), "score": 72 } }))
        .idempotency_key(random_uuid()), // safe to retry: the same key never writes twice
)?;
let ada_id = ada["id"].as_str().unwrap_or_default().to_string();
println!("Created {} version {}", ada_id, ada["version"]);

let record = api.call(
    Method::GET,
    &format!("/records/{ada_id}"),
    Request::default(),
)?;
let data = &record["data"];
println!(
    "{} {} version {}",
    data["name"].as_str().unwrap_or(""),
    data["email"].as_str().unwrap_or(""),
    record["version"]
);

/// Read, change, write with the current version. On VERSION_CONFLICT: read again and retry.
fn update_record(
    api: &Api,
    id: &str,
    change: impl Fn(&Value) -> Value,
) -> Result<Value, Box<dyn Error>> {
    let attempts = 3;
    for i in 1..=attempts {
        let current = api.call(Method::GET, &format!("/records/{id}"), Request::default())?;
        let body = json!({ "version": current["version"], "data": change(&current["data"]) });
        match api.call(
            Method::PATCH,
            &format!("/records/{id}"),
            Request::json(body),
        ) {
            Err(err)
                if i < attempts
                    && err
                        .downcast_ref::<AgentDBError>()
                        .is_some_and(|e| e.code == "VERSION_CONFLICT") =>
            {
                continue
            }
            result => return result,
        }
    }
    unreachable!("the last attempt always returns")
}

let updated = update_record(
    &api,
    &ada_id,
    |data| json!({ "score": data["score"].as_i64().unwrap_or(0) + 10 }),
)?;
println!(
    "Score {} version {}",
    updated["data"]["score"], updated["version"]
);

let stale = json!({ "version": 1, "data": { "score": 0 } }); // stale version
if let Err(err) = api.call(
    Method::PATCH,
    &format!("/records/{ada_id}"),
    Request::json(stale),
) {
    match err.downcast_ref::<AgentDBError>() {
        Some(e) if e.code == "VERSION_CONFLICT" => println!(
            "Conflict: the record is now at version {}",
            e.details["current_version"]
        ),
        _ => return Err(err),
    }
}

/// Every record that matches, page by page (limit up to 200; follow next_cursor).
fn query_all(
    api: &Api,
    object: &str,
    filter: &Value,
    sort: &Value,
    page_size: u32,
) -> Result<Vec<Value>, Box<dyn Error>> {
    let mut out = Vec::new();
    let mut cursor: Option<String> = None;
    loop {
        let mut query =
            json!({ "object": object, "where": filter, "sort": sort, "limit": page_size });
        if let Some(cursor) = &cursor {
            query["cursor"] = json!(cursor);
        }
        let mut page = api.call(Method::POST, "/query", Request::json(query))?;
        if let Value::Array(records) = page["records"].take() {
            out.extend(records);
        }
        cursor = page["next_cursor"].as_str().map(String::from); // null on the last page
        if cursor.is_none() {
            return Ok(out);
        }
    }
}

let good = query_all(
    &api,
    "contact",
    &json!({ "all": [{ "field": "score", "op": "gte", "value": 70 }, { "field": "email", "op": "neq", "value": "nobody@example.com" }] }),
    &json!([{ "field": "score", "order": "desc" }]),
    2,
)?;
for c in &good {
    println!(
        "{} {}",
        c["data"]["score"],
        c["data"]["name"].as_str().unwrap_or("")
    );
}

let company = api.call(
    Method::POST,
    "/records",
    Request::json(
        json!({ "object": "company", "data": { "name": format!("Example Ltd {run}") } }),
    ),
)?;
let company_id = company["id"].as_str().unwrap_or_default();
api.call(
    Method::POST,
    "/links",
    Request::json(json!({ "relationship": "works_at", "source_record_id": ada_id, "target_record_id": company_id }))
        .idempotency_key(format!("link-{ada_id}-{company_id}")),
)?;
let related = api.call(
    Method::GET,
    &format!("/records/{ada_id}/related?relationship=works_at"),
    Request::default(),
)?;
let employers: Vec<&str> = related["records"]
    .as_array()
    .into_iter()
    .flatten()
    .filter_map(|r| r["data"]["name"].as_str())
    .collect();
println!("Ada works at {}", employers.join(", "));

let bytes = format!("Meeting notes for Ada ({run})\n").into_bytes();
// 1. Announce the file, 2. send the raw bytes, 3. attach it to a record.
let file = api.call(
    Method::POST,
    "/files",
    Request::json(
        json!({ "filename": "notes.txt", "content_type": "text/plain", "size": bytes.len() }),
    ),
)?;
let file_id = file["id"].as_str().unwrap_or_default();
api.call(
    Method::PUT,
    &format!("/files/{file_id}/content"),
    Request {
        body: Some(bytes.clone()),
        ..Default::default()
    },
)?;
api.call(
    Method::POST,
    "/attachments",
    Request::json(json!({ "file_id": file_id, "record_id": ada_id })),
)?;
// Download with the key (server side) ...
let back = api
    .send(
        Method::GET,
        &format!("/files/{file_id}/content"),
        &Request::default(),
    )?
    .bytes()?;
// ... or make a signed link that expires, for an <img> or a download button in your page.
let link = api.call(
    Method::POST,
    &format!("/files/{file_id}/links"),
    Request::json(json!({ "expires_in_seconds": 600 })),
)?;
println!(
    "Signed link until {}",
    link["expires_at"].as_str().unwrap_or("")
);

if let Err(err) = api.call(Method::GET, "/records/rec_doesnotexist", Request::default()) {
    let Some(e) = err.downcast_ref::<AgentDBError>() else {
        return Err(err); // a network error, not an AgentDB answer
    };
    println!("{} {} {}", e.status, e.code, e.message); // 404 NOT_FOUND ...
    if e.retryable {
        println!("Safe to retry later");
    }
}

let key = random_uuid(); // one key per intention; reuse it only to retry the same request
let request = json!({ "object": "contact", "data": { "name": "Mary Somerville", "email": format!("mary-{run}@example.com") } });
let first = api.call(
    Method::POST,
    "/records",
    Request::json(request.clone()).idempotency_key(&key),
)?;
let retry = api.call(
    Method::POST,
    "/records",
    Request::json(request).idempotency_key(&key),
)?; // e.g. after a timeout
println!(
    "{}",
    if first["id"] == retry["id"] {
        "Same record, written once"
    } else {
        "Duplicate!"
    }
);

cargo run