roundabout

public

Star0

Watch0

11 commits 1 branch 0 tags 0 releases158.8 KiB

use std::{io, os::unix::net::UnixStream, path::PathBuf};

use pound::Parse;

use crate::{
    artifacts, client,
    config::Config,
    daemon, history, hook, paths,
    proto::{self, Event, Reader, Request},
    run,
    units::{Amount, Bytes, Dur},
};

/// one build at a time, please
#[derive(Parse)]
#[pound(name = "rabt", version = env!("CARGO_PKG_VERSION"))]
pub enum Rabt {
    /// wait your turn, then run the thing
    Run {
        /// preset from config, s m l xl out of the box
        #[pound(short, long)]
        size: Option<String>,
        #[pound(long)]
        ram: Option<Bytes>,
        #[pound(long)]
        disk: Option<Bytes>,
        /// count or percentage of the machine
        #[pound(long)]
        cores: Option<Amount<u32>>,
        /// how long you reckon it'll take
        #[pound(short, long)]
        time: Option<Dur>,
        /// kill it past this, exit 124
        #[pound(long)]
        deadline: Option<Dur>,
        /// stop queueing past this, exit 75
        #[pound(long)]
        max_wait: Option<Dur>,
        #[pound(long, hidden)]
        session: Option<String>,
        /// exactly the cores asked for, no spares
        #[pound(long)]
        exact: bool,
        #[pound(trailing)]
        cmd: Vec<String>,
    },
    /// say you'll be building soon, so others leave you room
    Expect {
        #[pound(short, long)]
        size: Option<String>,
        #[pound(long)]
        ram: Option<Bytes>,
        #[pound(long)]
        cores: Option<Amount<u32>>,
        /// changed your mind
        #[pound(long)]
        done: bool,
        #[pound(long, hidden)]
        session: Option<String>,
    },
    /// who's running, who's waiting
    Status {
        #[pound(long)]
        json: bool,
    },
    /// a size guess from past runs
    Suggest {
        #[pound(trailing)]
        cmd: Vec<String>,
    },
    /// stop a lease's build
    Kill { lease: u64 },
    /// buy an overdue build more time
    Extend { lease: u64, by: Dur },
    /// exit when a lease goes overdue or finishes
    Watch { lease: u64 },
    /// build output and how much of it there is
    Artifacts {
        #[pound(subcommand)]
        action: Artifacts,
    },
    /// agent hook entrypoint, json on stdin
    Hook {
        /// pre-tool-use post-tool-use prompt-submit stop subagent-stop session-end, or codex-prefixed variants
        event: String,
    },
    #[pound(hidden)]
    Daemon,
}

#[derive(Parse)]
pub enum Artifacts {
    /// sizes against the caps, exits 1 when over
    Check {
        /// re-measure instead of trusting the cache
        #[pound(long)]
        fresh: bool,
    },
    /// everything tracked, biggest first
    List,
    /// delete idle artifacts
    Gc {
        /// keep going, oldest first, until under the caps
        #[pound(long)]
        cap: bool,
        #[pound(short = 'n', long)]
        dry_run: bool,
    },
    /// never gc this
    Pin {
        path: PathBuf,
    },
    Unpin {
        path: PathBuf,
    },
}

fn artifacts(cfg: &Config, action: Artifacts) -> i32 {
    let cwd = std::env::current_dir().unwrap_or_else(|_| "/".into());
    let root = paths::project_root(&cwd);
    match action {
        Artifacts::Check { fresh } => {
            let r = artifacts::check(cfg, Some(&root), fresh);
            if let Some((root, size)) = &r.project {
                println!(
                    "{}: {size} of {}",
                    root.display(),
                    cfg.artifacts.project_cap
                );
            }
            println!("everything: {} of {}", r.global, cfg.artifacts.global_cap);
            if r.over() {
                println!("over. `rabt artifacts gc --cap`");
                1
            } else {
                0
            }
        }
        Artifacts::List => {
            let r = artifacts::check(cfg, Some(&root), false);
            let mut rows: Vec<_> = r.entries.iter().collect();
            rows.sort_by_key(|(_, e)| std::cmp::Reverse(e.size));
            for (path, e) in rows {
                println!(
                    "{:>10}  {}{}",
                    e.size.to_string(),
                    path.display(),
                    if e.pinned { "  (pinned)" } else { "" }
                );
            }
            0
        }
        Artifacts::Gc { cap, dry_run } => {
            let g = artifacts::gc(cfg, cap, dry_run);
            for p in &g.removed {
                println!(
                    "{}{}",
                    if dry_run { "would remove " } else { "removed " },
                    p.display()
                );
            }
            for f in &g.failed {
                eprintln!("rabt: {f}");
            }
            println!(
                "{} {}",
                if dry_run { "would free" } else { "freed" },
                g.freed
            );
            i32::from(!g.failed.is_empty())
        }
        Artifacts::Pin { path } | Artifacts::Unpin { path } if !path.exists() => {
            eprintln!("rabt: {} doesn't exist", path.display());
            1
        }
        Artifacts::Pin { path } => report(artifacts::pin(&path, true)),
        Artifacts::Unpin { path } => report(artifacts::pin(&path, false)),
    }
}

fn report(r: Result<(), String>) -> i32 {
    r.map_or_else(
        |e| {
            eprintln!("rabt: {e}");
            1
        },
        |()| 0,
    )
}

fn ask(req: &Request) -> io::Result<Option<Event>> {
    let Some(mut s) = client::connect() else {
        return Ok(None);
    };
    proto::send(&mut s, req)?;
    Reader::new(&s)?.recv()
}

fn reply(ev: io::Result<Option<Event>>) -> i32 {
    match ev {
        Ok(Some(Event::Ok)) => 0,
        Ok(Some(Event::Err { msg })) => {
            eprintln!("rabt: {msg}");
            1
        }
        Ok(None) => {
            eprintln!("rabt: nothing running");
            1
        }
        Ok(Some(_)) => 1,
        Err(e) => {
            eprintln!("rabt: {e}");
            1
        }
    }
}

fn agent_session(session: Option<String>) -> Option<String> {
    session.or_else(|| {
        [
            "ROUNDABOUT_SESSION_ID",
            "CLAUDE_CODE_SESSION_ID",
            "CODEX_THREAD_ID",
            "CODEX_SESSION_ID",
        ]
        .into_iter()
        .find_map(|key| std::env::var(key).ok().filter(|value| !value.is_empty()))
    })
}

fn expect(cfg: &Config, session: Option<String>, done: bool, a: &run::Args) -> i32 {
    let Some(session) = session.filter(|s| !s.is_empty()) else {
        eprintln!("rabt: no agent session id, use --session to name one");
        return 2;
    };
    let need = if done {
        None
    } else {
        match run::need(cfg, a) {
            Ok(n) => Some(n),
            Err(e) => {
                eprintln!("rabt: {e}");
                return 2;
            }
        }
    };
    let mut s = match client::connect_or_spawn() {
        Ok(s) => s,
        Err(e) => {
            eprintln!("rabt: can't reach the daemon: {e}");
            return 1;
        }
    };
    if let Err(e) = proto::send(&mut s, &Request::Expect { session, need }) {
        eprintln!("rabt: {e}");
        return 1;
    }
    reply(Reader::new(&s).and_then(|mut r| r.recv()))
}

fn status(json: bool) -> i32 {
    let Ok(Some(Event::Status(st))) = ask(&Request::Status) else {
        if json {
            println!("null");
        } else {
            println!("idle");
        }
        return 0;
    };
    if json {
        println!("{}", serde_json::to_string(&st).unwrap_or_default());
        return 0;
    }
    println!(
        "ram {} free of {}, cores {} free of {}",
        st.ram_free, st.ram_budget, st.cores_free, st.cores
    );
    for l in &st.leases {
        let state = if l.running {
            format!(
                "running {}{}",
                Dur::secs(l.elapsed),
                if l.overdue { " OVERDUE" } else { "" }
            )
        } else {
            format!(
                "#{} {}",
                l.position,
                l.waiting.map(|r| r.to_string()).unwrap_or_default()
            )
        };
        let cores = if l.granted > l.need.cores {
            format!("{}+{}", l.need.cores, l.granted - l.need.cores)
        } else {
            l.need.cores.to_string()
        };
        println!(
            "{:>4} {}{} {} ram {} cores {} rss {}  {}",
            l.id,
            if l.heavy { "H " } else { "" },
            state,
            Dur(l.need.time),
            l.need.ram,
            cores,
            l.rss,
            l.cmd.join(" ")
        );
    }
    for i in &st.intents {
        println!(
            "     expect {} ram {} cores {}, {} left",
            &i.session[..i.session.len().min(8)],
            i.need.ram,
            i.need.cores,
            Dur::secs(i.expires_in)
        );
    }
    0
}

fn suggest(cfg: &Config, cmd: &[String]) -> i32 {
    let cwd = std::env::current_dir().unwrap_or_else(|_| "/".into());
    let key = history::key(&paths::project_root(&cwd), cmd);
    let Some(rec) = history::lookup(&key) else {
        println!("never seen it, pick a --size");
        return 1;
    };
    let base = match run::base(cfg) {
        Ok(b) => b,
        Err(e) => {
            eprintln!("rabt: {e}");
            return 2;
        }
    };
    let n = history::suggest(&rec, base, cmd);
    println!(
        "--ram {} --cores {} --time {} --disk {}",
        n.ram,
        n.cores,
        Dur(n.time),
        n.disk
    );
    0
}

fn watch(lease: u64) -> i32 {
    let Some(mut s): Option<UnixStream> = client::connect() else {
        return 0;
    };
    if proto::send(&mut s, &Request::Watch { lease }).is_err() {
        return 0;
    }
    if matches!(
        Reader::new(&s).and_then(|mut r| r.recv::<Event>()),
        Ok(Some(Event::Fired { overdue: true }))
    ) {
        println!("lease {lease} is overdue, `rabt kill {lease}` or `rabt extend {lease} <dur>`");
        3
    } else {
        println!("lease {lease} is done");
        0
    }
}

pub fn dispatch(cfg: &Config, cmd: Rabt) -> i32 {
    match cmd {
        Rabt::Run {
            size,
            ram,
            disk,
            cores,
            time,
            deadline,
            max_wait,
            session,
            exact,
            cmd,
        } => run::run(
            cfg,
            &run::Args {
                size,
                ram,
                disk,
                cores,
                time,
                deadline,
                max_wait,
                session: agent_session(session),
                exact,
                cmd,
            },
        ),
        Rabt::Expect {
            size,
            ram,
            cores,
            done,
            session,
        } => expect(
            cfg,
            agent_session(session),
            done,
            &run::Args {
                size,
                ram,
                cores,
                ..run::Args::default()
            },
        ),
        Rabt::Status { json } => status(json),
        Rabt::Suggest { cmd } => suggest(cfg, &cmd),
        Rabt::Kill { lease } => reply(ask(&Request::Kill { lease })),
        Rabt::Extend { lease, by } => reply(ask(&Request::Extend { lease, by })),
        Rabt::Watch { lease } => watch(lease),
        Rabt::Artifacts { action } => artifacts(cfg, action),
        Rabt::Hook { event } => hook::main(cfg, &event),
        Rabt::Daemon => daemon::main(Config::load().unwrap_or_default()),
    }
}

1	use std::{io, os::unix::net::UnixStream, path::PathBuf};
2
3	use pound::Parse;
4
5	use crate::{
6	artifacts, client,
7	config::Config,
8	daemon, history, hook, paths,
9	proto::{self, Event, Reader, Request},
10	run,
11	units::{Amount, Bytes, Dur},
12	};
13
14	/// one build at a time, please
15	#[derive(Parse)]
16	#[pound(name = "rabt", version = env!("CARGO_PKG_VERSION"))]
17	pub enum Rabt {
18	/// wait your turn, then run the thing
19	Run {
20	/// preset from config, s m l xl out of the box
21	#[pound(short, long)]
22	size: Option<String>,
23	#[pound(long)]
24	ram: Option<Bytes>,
25	#[pound(long)]
26	disk: Option<Bytes>,
27	/// count or percentage of the machine
28	#[pound(long)]
29	cores: Option<Amount<u32>>,
30	/// how long you reckon it'll take
31	#[pound(short, long)]
32	time: Option<Dur>,
33	/// kill it past this, exit 124
34	#[pound(long)]
35	deadline: Option<Dur>,
36	/// stop queueing past this, exit 75
37	#[pound(long)]
38	max_wait: Option<Dur>,
39	#[pound(long, hidden)]
40	session: Option<String>,
41	/// exactly the cores asked for, no spares
42	#[pound(long)]
43	exact: bool,
44	#[pound(trailing)]
45	cmd: Vec<String>,
46	},
47	/// say you'll be building soon, so others leave you room
48	Expect {
49	#[pound(short, long)]
50	size: Option<String>,
51	#[pound(long)]
52	ram: Option<Bytes>,
53	#[pound(long)]
54	cores: Option<Amount<u32>>,
55	/// changed your mind
56	#[pound(long)]
57	done: bool,
58	#[pound(long, hidden)]
59	session: Option<String>,
60	},
61	/// who's running, who's waiting
62	Status {
63	#[pound(long)]
64	json: bool,
65	},
66	/// a size guess from past runs
67	Suggest {
68	#[pound(trailing)]
69	cmd: Vec<String>,
70	},
71	/// stop a lease's build
72	Kill { lease: u64 },
73	/// buy an overdue build more time
74	Extend { lease: u64, by: Dur },
75	/// exit when a lease goes overdue or finishes
76	Watch { lease: u64 },
77	/// build output and how much of it there is
78	Artifacts {
79	#[pound(subcommand)]
80	action: Artifacts,
81	},
82	/// agent hook entrypoint, json on stdin
83	Hook {
84	/// pre-tool-use post-tool-use prompt-submit stop subagent-stop session-end, or codex-prefixed variants
85	event: String,
86	},
87	#[pound(hidden)]
88	Daemon,
89	}
90
91	#[derive(Parse)]
92	pub enum Artifacts {
93	/// sizes against the caps, exits 1 when over
94	Check {
95	/// re-measure instead of trusting the cache
96	#[pound(long)]
97	fresh: bool,
98	},
99	/// everything tracked, biggest first
100	List,
101	/// delete idle artifacts
102	Gc {
103	/// keep going, oldest first, until under the caps
104	#[pound(long)]
105	cap: bool,
106	#[pound(short = 'n', long)]
107	dry_run: bool,
108	},
109	/// never gc this
110	Pin {
111	path: PathBuf,
112	},
113	Unpin {
114	path: PathBuf,
115	},
116	}
117
118	fn artifacts(cfg: &Config, action: Artifacts) -> i32 {
119	let cwd = std::env::current_dir().unwrap_or_else(\|_\| "/".into());
120	let root = paths::project_root(&cwd);
121	match action {
122	Artifacts::Check { fresh } => {
123	let r = artifacts::check(cfg, Some(&root), fresh);
124	if let Some((root, size)) = &r.project {
125	println!(
126	"{}: {size} of {}",
127	root.display(),
128	cfg.artifacts.project_cap
129	);
130	}
131	println!("everything: {} of {}", r.global, cfg.artifacts.global_cap);
132	if r.over() {
133	println!("over. `rabt artifacts gc --cap`");
134	1
135	} else {
136	0
137	}
138	}
139	Artifacts::List => {
140	let r = artifacts::check(cfg, Some(&root), false);
141	let mut rows: Vec<_> = r.entries.iter().collect();
142	rows.sort_by_key(\|(_, e)\| std::cmp::Reverse(e.size));
143	for (path, e) in rows {
144	println!(
145	"{:>10} {}{}",
146	e.size.to_string(),
147	path.display(),
148	if e.pinned { " (pinned)" } else { "" }
149	);
150	}
151	0
152	}
153	Artifacts::Gc { cap, dry_run } => {
154	let g = artifacts::gc(cfg, cap, dry_run);
155	for p in &g.removed {
156	println!(
157	"{}{}",
158	if dry_run { "would remove " } else { "removed " },
159	p.display()
160	);
161	}
162	for f in &g.failed {
163	eprintln!("rabt: {f}");
164	}
165	println!(
166	"{} {}",
167	if dry_run { "would free" } else { "freed" },
168	g.freed
169	);
170	i32::from(!g.failed.is_empty())
171	}
172	Artifacts::Pin { path } \| Artifacts::Unpin { path } if !path.exists() => {
173	eprintln!("rabt: {} doesn't exist", path.display());
174	1
175	}
176	Artifacts::Pin { path } => report(artifacts::pin(&path, true)),
177	Artifacts::Unpin { path } => report(artifacts::pin(&path, false)),
178	}
179	}
180
181	fn report(r: Result<(), String>) -> i32 {
182	r.map_or_else(
183	\|e\| {
184	eprintln!("rabt: {e}");
185	1
186	},
187	\|()\| 0,
188	)
189	}
190
191	fn ask(req: &Request) -> io::Result<Option<Event>> {
192	let Some(mut s) = client::connect() else {
193	return Ok(None);
194	};
195	proto::send(&mut s, req)?;
196	Reader::new(&s)?.recv()
197	}
198
199	fn reply(ev: io::Result<Option<Event>>) -> i32 {
200	match ev {
201	Ok(Some(Event::Ok)) => 0,
202	Ok(Some(Event::Err { msg })) => {
203	eprintln!("rabt: {msg}");
204	1
205	}
206	Ok(None) => {
207	eprintln!("rabt: nothing running");
208	1
209	}
210	Ok(Some(_)) => 1,
211	Err(e) => {
212	eprintln!("rabt: {e}");
213	1
214	}
215	}
216	}
217
218	fn agent_session(session: Option<String>) -> Option<String> {
219	session.or_else(\|\| {
220	[
221	"ROUNDABOUT_SESSION_ID",
222	"CLAUDE_CODE_SESSION_ID",
223	"CODEX_THREAD_ID",
224	"CODEX_SESSION_ID",
225	]
226	.into_iter()
227	.find_map(\|key\| std::env::var(key).ok().filter(\|value\| !value.is_empty()))
228	})
229	}
230
231	fn expect(cfg: &Config, session: Option<String>, done: bool, a: &run::Args) -> i32 {
232	let Some(session) = session.filter(\|s\| !s.is_empty()) else {
233	eprintln!("rabt: no agent session id, use --session to name one");
234	return 2;
235	};
236	let need = if done {
237	None
238	} else {
239	match run::need(cfg, a) {
240	Ok(n) => Some(n),
241	Err(e) => {
242	eprintln!("rabt: {e}");
243	return 2;
244	}
245	}
246	};
247	let mut s = match client::connect_or_spawn() {
248	Ok(s) => s,
249	Err(e) => {
250	eprintln!("rabt: can't reach the daemon: {e}");
251	return 1;
252	}
253	};
254	if let Err(e) = proto::send(&mut s, &Request::Expect { session, need }) {
255	eprintln!("rabt: {e}");
256	return 1;
257	}
258	reply(Reader::new(&s).and_then(\|mut r\| r.recv()))
259	}
260
261	fn status(json: bool) -> i32 {
262	let Ok(Some(Event::Status(st))) = ask(&Request::Status) else {
263	if json {
264	println!("null");
265	} else {
266	println!("idle");
267	}
268	return 0;
269	};
270	if json {
271	println!("{}", serde_json::to_string(&st).unwrap_or_default());
272	return 0;
273	}
274	println!(
275	"ram {} free of {}, cores {} free of {}",
276	st.ram_free, st.ram_budget, st.cores_free, st.cores
277	);
278	for l in &st.leases {
279	let state = if l.running {
280	format!(
281	"running {}{}",
282	Dur::secs(l.elapsed),
283	if l.overdue { " OVERDUE" } else { "" }
284	)
285	} else {
286	format!(
287	"#{} {}",
288	l.position,
289	l.waiting.map(\|r\| r.to_string()).unwrap_or_default()
290	)
291	};
292	let cores = if l.granted > l.need.cores {
293	format!("{}+{}", l.need.cores, l.granted - l.need.cores)
294	} else {
295	l.need.cores.to_string()
296	};
297	println!(
298	"{:>4} {}{} {} ram {} cores {} rss {} {}",
299	l.id,
300	if l.heavy { "H " } else { "" },
301	state,
302	Dur(l.need.time),
303	l.need.ram,
304	cores,
305	l.rss,
306	l.cmd.join(" ")
307	);
308	}
309	for i in &st.intents {
310	println!(
311	" expect {} ram {} cores {}, {} left",
312	&i.session[..i.session.len().min(8)],
313	i.need.ram,
314	i.need.cores,
315	Dur::secs(i.expires_in)
316	);
317	}
318	0
319	}
320
321	fn suggest(cfg: &Config, cmd: &[String]) -> i32 {
322	let cwd = std::env::current_dir().unwrap_or_else(\|_\| "/".into());
323	let key = history::key(&paths::project_root(&cwd), cmd);
324	let Some(rec) = history::lookup(&key) else {
325	println!("never seen it, pick a --size");
326	return 1;
327	};
328	let base = match run::base(cfg) {
329	Ok(b) => b,
330	Err(e) => {
331	eprintln!("rabt: {e}");
332	return 2;
333	}
334	};
335	let n = history::suggest(&rec, base, cmd);
336	println!(
337	"--ram {} --cores {} --time {} --disk {}",
338	n.ram,
339	n.cores,
340	Dur(n.time),
341	n.disk
342	);
343	0
344	}
345
346	fn watch(lease: u64) -> i32 {
347	let Some(mut s): Option<UnixStream> = client::connect() else {
348	return 0;
349	};
350	if proto::send(&mut s, &Request::Watch { lease }).is_err() {
351	return 0;
352	}
353	if matches!(
354	Reader::new(&s).and_then(\|mut r\| r.recv::<Event>()),
355	Ok(Some(Event::Fired { overdue: true }))
356	) {
357	println!("lease {lease} is overdue, `rabt kill {lease}` or `rabt extend {lease} <dur>`");
358	3
359	} else {
360	println!("lease {lease} is done");
361	0
362	}
363	}
364
365	pub fn dispatch(cfg: &Config, cmd: Rabt) -> i32 {
366	match cmd {
367	Rabt::Run {
368	size,
369	ram,
370	disk,
371	cores,
372	time,
373	deadline,
374	max_wait,
375	session,
376	exact,
377	cmd,
378	} => run::run(
379	cfg,
380	&run::Args {
381	size,
382	ram,
383	disk,
384	cores,
385	time,
386	deadline,
387	max_wait,
388	session: agent_session(session),
389	exact,
390	cmd,
391	},
392	),
393	Rabt::Expect {
394	size,
395	ram,
396	cores,
397	done,
398	session,
399	} => expect(
400	cfg,
401	agent_session(session),
402	done,
403	&run::Args {
404	size,
405	ram,
406	cores,
407	..run::Args::default()
408	},
409	),
410	Rabt::Status { json } => status(json),
411	Rabt::Suggest { cmd } => suggest(cfg, &cmd),
412	Rabt::Kill { lease } => reply(ask(&Request::Kill { lease })),
413	Rabt::Extend { lease, by } => reply(ask(&Request::Extend { lease, by })),
414	Rabt::Watch { lease } => watch(lease),
415	Rabt::Artifacts { action } => artifacts(cfg, action),
416	Rabt::Hook { event } => hook::main(cfg, &event),
417	Rabt::Daemon => daemon::main(Config::load().unwrap_or_default()),
418	}
419	}