use crate::agent::client::AgentClient;
use crate::agent::hub::{self, HubClient, HubError};
use std::path::PathBuf;
use std::time::Duration;
const T: Duration = Duration::from_secs(5);
pub fn explicit_workers(args: &[String]) -> Option<u32> {
let mut it = args.iter();
while let Some(a) = it.next() {
match a.as_str() {
"--workers" | "-w" | "-n" => return it.next().and_then(|v| v.parse().ok()),
s if s.starts_with("--workers=") => return s["--workers=".len()..].parse().ok(),
_ => {}
}
}
None
}
pub fn port_flags_given(args: &[String]) -> bool {
args.iter()
.any(|a| matches!(a.as_str(), "--port" | "-p" | "--broker-port" | "-b"))
}
fn num(v: f64) -> String {
format!("{v}")
}
pub fn format_price(state: &str, per_hour: Option<f64>, effective: Option<f64>) -> String {
match (state, per_hour, effective) {
("set", Some(p), _) => format!("this Mac: {} credits/hour", num(p)),
("inherited", _, Some(e)) => format!("this Mac: account default ({} credits/hour)", num(e)),
("inherited", _, None) => "this Mac: account default".to_string(),
("mixed", _, Some(e)) => format!("this Mac: mixed (from {} credits/hour)", num(e)),
("disabled", _, _) => "this Mac: disabled on the hub".to_string(),
_ => "this Mac: no price yet".to_string(),
}
}
fn print_error(e: String) -> i32 {
eprintln!("zc: {e}");
1
}
pub fn up(c: &AgentClient, args: &[String]) -> i32 {
if port_flags_given(args) {
eprintln!(" note: the zc agent owns the ports; --port/--broker-port are ignored");
}
if let Some(n) = explicit_workers(args) {
if let Err(e) = c.put("/v1/workers", serde_json::json!({ "count": n }), T) {
return print_error(e);
}
}
match c.put("/v1/sharing", serde_json::json!({ "on": true }), T) {
Ok(s) => {
println!(
"✓ Sharing handed to zc agent: {} worker(s). Follow it with `zc agent status`.",
s["this_mac"]["workers"]["desired"]
);
0
}
Err(e) => print_error(e),
}
}
pub fn down(c: &AgentClient) -> Option<i32> {
match c.put("/v1/sharing", serde_json::json!({ "on": false }), T) {
Ok(s) if s["this_mac"]["broker"] == "unmanaged" => {
println!("✓ zc agent paused sharing; stopping the broker it didn't start.");
None
}
Ok(_) => {
println!("✓ zc agent is draining the workers and will stop sharing.");
Some(0)
}
Err(e) => Some(print_error(e)),
}
}
fn finite_price(v: f64) -> Result<f64, String> {
if v.is_finite() {
Ok(v)
} else {
Err(format!(
"price must be a finite number of credits/hour, got {v}"
))
}
}
pub fn price_set(c: &AgentClient, value: f64) -> i32 {
if let Err(e) = finite_price(value) {
return print_error(e);
}
let body = serde_json::json!({ "scope": "device", "price_per_hour": value });
match c.put("/v1/price", body, Duration::from_secs(30)) {
Ok(_) => {
println!("price set: {} credits/hour on this Mac (hub)", num(value));
0
}
Err(e) => print_error(e),
}
}
pub fn price_show(c: &AgentClient) -> i32 {
match c.summary(T) {
Ok(s) => {
let p = &s["prices"]["this_mac"];
println!(
"{}",
format_price(
p["state"].as_str().unwrap_or(""),
p["per_hour"].as_f64(),
p["effective_per_hour"].as_f64()
)
);
0
}
Err(e) => print_error(e),
}
}
pub fn try_up(args: &[String]) -> Option<i32> {
AgentClient::running().map(|c| up(&c, args))
}
pub fn try_down() -> Option<i32> {
AgentClient::running().and_then(|c| down(&c))
}
pub fn try_price(value: Option<f64>) -> Option<i32> {
let c = AgentClient::running()?;
Some(match value {
Some(v) => price_set(&c, v),
None => price_show(&c),
})
}
async fn summary_with_retry(
client: &HubClient,
tz: &str,
) -> Result<crate::agent::hubwire::HubSummary, HubError> {
let result = client.summary(tz).await;
if tz != "UTC" && matches!(result, Err(HubError::Status(422, _))) {
return client.summary("UTC").await;
}
result
}
pub fn price_direct(value: Option<f64>) -> i32 {
price_direct_in(crate::credentials::dir(), value)
}
fn price_direct_in(state_dir: Option<PathBuf>, value: Option<f64>) -> i32 {
if let Some(Err(e)) = value.map(finite_price) {
return print_error(e);
}
let Some(state_dir) = state_dir else {
return print_error("no HOME or ZAKURO_HOME".into());
};
let Some(creds) = hub::read_credentials(&state_dir.join("credentials")) else {
return print_error("this Mac is not signed in: run `zc login`".into());
};
let base = creds
.api_url
.clone()
.unwrap_or_else(crate::credentials::default_api_url);
let pk = crate::broker::node_identity::NodeKey::load_or_create_in(Some(state_dir)).public_b64();
let rt = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(rt) => rt,
Err(e) => return print_error(e.to_string()),
};
rt.block_on(async move {
let client = HubClient::new(&base, &creds.api_key, Duration::from_secs(10));
let summary = match summary_with_retry(&client, &hub::system_tz()).await {
Ok(s) => s,
Err(e) => return print_error(e.message()),
};
let mine = summary.this_mac_workers(&pk);
let Some(v) = value else {
let p = crate::agent::merge::this_mac_price(&mine);
println!(
"{}",
format_price(&p.state, p.per_hour, p.effective_per_hour)
);
return 0;
};
let b = summary.bounds();
if v < b.min as f64 || v > b.max as f64 {
return print_error(format!(
"price must be between {} and {} credits/hour",
b.min, b.max
));
}
let ids: Vec<i64> = mine.iter().map(|w| w.id).collect();
if ids.is_empty() {
return print_error(
"this Mac has no workers on the hub yet: start sharing first (`zc share`)".into(),
);
}
match hub::set_device_price(&client, &ids, Some(v)).await {
Ok(()) => {
println!(
"price set: {} credits/hour on this Mac ({} worker(s))",
num(v),
ids.len()
);
0
}
Err(failed) => print_error(format!("could not price {} worker(s)", failed.len())),
}
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agent::client::tests::{idle_agent, sleeper_cfg, start_agent};
use crate::agent::core::tests::{hub_json, row};
use crate::agent::hub::tests::{mock_hub, Seen};
fn args(s: &[&str]) -> Vec<String> {
s.iter().map(|a| a.to_string()).collect()
}
#[test]
fn worker_and_port_flags_are_detected() {
assert_eq!(explicit_workers(&args(&["--workers", "3"])), Some(3));
assert_eq!(explicit_workers(&args(&["-w", "2", "-d"])), Some(2));
assert_eq!(explicit_workers(&args(&["-n", "4"])), Some(4));
assert_eq!(explicit_workers(&args(&["--workers=5"])), Some(5));
assert_eq!(explicit_workers(&args(&["-d"])), None);
assert!(port_flags_given(&args(&["--port", "4000"])));
assert!(port_flags_given(&args(&["-b", "9100"])));
assert!(!port_flags_given(&args(&["--workers", "2"])));
}
#[test]
fn price_lines() {
assert_eq!(
format_price("set", Some(18.0), Some(18.0)),
"this Mac: 18 credits/hour"
);
assert_eq!(
format_price("inherited", None, Some(12.0)),
"this Mac: account default (12 credits/hour)"
);
assert_eq!(
format_price("inherited", None, None),
"this Mac: account default"
);
assert_eq!(
format_price("mixed", None, Some(12.0)),
"this Mac: mixed (from 12 credits/hour)"
);
}
#[test]
fn a_disabled_price_state_is_shown_as_disabled_on_the_hub() {
assert_eq!(
format_price("disabled", None, None),
"this Mac: disabled on the hub"
);
}
#[test]
fn a_422_on_the_starting_zone_is_retried_once_in_utc() {
use crate::agent::hubwire::tests::HUB_JSON;
let (base, seen) = mock_hub(|_, url, _| match url {
"/api/provide/summary?tz=Asia%2FTokyo" => (422, "{}".into()),
"/api/provide/summary?tz=UTC" => (200, HUB_JSON.to_string()),
_ => (404, "{}".into()),
});
let client = hub::HubClient::new(&base, "zk_1_x", std::time::Duration::from_secs(5));
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let summary = rt
.block_on(summary_with_retry(&client, "Asia/Tokyo"))
.expect("retried in UTC and succeeded");
assert_eq!(summary.account.username, "jean");
let urls: Vec<String> = seen
.lock()
.unwrap()
.iter()
.map(|(_, u, _, _)| u.clone())
.collect();
assert_eq!(
urls,
vec![
"/api/provide/summary?tz=Asia%2FTokyo".to_string(),
"/api/provide/summary?tz=UTC".to_string()
]
);
}
#[test]
fn up_and_down_are_handed_to_the_agent() {
let (c, _stop) = idle_agent();
assert_eq!(up(&c, &args(&["--workers", "2", "--port", "4000"])), 0);
let s = c.summary(std::time::Duration::from_secs(2)).unwrap();
assert_eq!(
(
s["this_mac"]["sharing"].as_bool(),
s["this_mac"]["workers"]["desired"].as_u64()
),
(Some(true), Some(2))
);
assert_eq!(down(&c), Some(0));
assert_eq!(
c.summary(std::time::Duration::from_secs(2)).unwrap()["this_mac"]["sharing"],
false
);
assert_eq!(
price_set(&c, 10.0),
1,
"not signed in: the agent answers 503"
);
}
fn wait_until(what: &str, mut done: impl FnMut() -> bool) {
let end = std::time::Instant::now() + Duration::from_secs(10);
while std::time::Instant::now() < end {
if done() {
return;
}
std::thread::sleep(Duration::from_millis(100));
}
panic!("timed out waiting for: {what}");
}
#[test]
fn down_with_an_unmanaged_broker_pauses_the_agent_then_falls_through() {
let (fake_broker, _) = mock_hub(|_, url, _| match url {
"/health" => (
200,
r#"{"status":"healthy","service":"zakuro-broker"}"#.into(),
),
_ => (404, "{}".into()),
});
let mut cfg = sleeper_cfg("down-unmanaged");
cfg.broker_port = fake_broker.rsplit(':').next().unwrap().parse().unwrap();
crate::agent::files::ensure_dir(&cfg.dir).unwrap();
crate::agent::files::save_json(
&cfg.dir.join(crate::agent::files::STATE_FILE),
&crate::agent::files::DesiredState {
sharing: true,
..Default::default()
},
)
.unwrap();
let (c, _guard) = start_agent(cfg);
let summary = || c.summary(Duration::from_secs(2)).unwrap();
wait_until("the agent sees the unmanaged broker", || {
summary()["this_mac"]["broker"] == "unmanaged"
});
assert_eq!(summary()["this_mac"]["sharing"], true);
assert_eq!(down(&c), None, "the port-based `zc down` must still run");
assert_eq!(
summary()["this_mac"]["sharing"],
false,
"the agent paused sharing first"
);
}
fn pricing_hub(pk: String) -> (String, Seen) {
mock_hub(move |method, url, _| match method {
"GET" if url.starts_with("/api/provide/summary") => {
(200, hub_json(&pk, &row(11, None)))
}
"PUT" => (200, "{}".into()),
_ => (404, "{}".into()),
})
}
fn put_bodies(seen: &Seen) -> Vec<String> {
seen.lock()
.unwrap()
.iter()
.filter(|e| e.0 == "PUT")
.map(|e| e.2.clone())
.collect()
}
#[test]
fn a_non_finite_price_is_refused_before_it_reaches_the_hub() {
let state = crate::agent::files::tests::tmp("price-direct");
std::fs::create_dir_all(&state).unwrap();
let pk = crate::broker::node_identity::NodeKey::load_or_create_in(Some(state.clone()))
.public_b64();
let (hub, seen) = pricing_hub(pk);
std::fs::write(
state.join("credentials"),
format!("api_key=zk_1_test\napi_url={hub}\n"),
)
.unwrap();
for v in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
assert_eq!(price_direct_in(Some(state.clone()), Some(v)), 1, "{v}");
}
assert_eq!(put_bodies(&seen), Vec::<String>::new(), "nothing sent");
assert_eq!(
price_direct_in(Some(state), Some(18.0)),
0,
"a real price still goes through"
);
assert_eq!(
put_bodies(&seen),
vec![r#"{"price_per_hour":18.0}"#.to_string()]
);
}
#[test]
fn a_non_finite_price_is_refused_before_it_reaches_the_agent() {
let cfg = sleeper_cfg("price-set");
std::fs::create_dir_all(&cfg.state_dir).unwrap();
let (hub, seen) = pricing_hub(cfg.node_pubkey());
std::fs::write(
cfg.state_dir.join("credentials"),
format!("api_key=zk_1_test\napi_url={hub}\n"),
)
.unwrap();
let (c, _guard) = start_agent(cfg);
for v in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
assert_eq!(price_set(&c, v), 1, "{v}");
}
assert_eq!(put_bodies(&seen), Vec::<String>::new(), "nothing sent");
assert_eq!(price_set(&c, 18.0), 0, "a real price still goes through");
assert_eq!(
put_bodies(&seen),
vec![r#"{"price_per_hour":18.0}"#.to_string()]
);
}
}