roundabout

public

Star0

Watch0

14 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 = "roundabout", 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,
        /// leave other projects alone
        #[pound(long)]
        project: 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. `{}`", r.remedy());
                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,
            project,
            dry_run,
        } => {
            let g = artifacts::gc(cfg, project.then_some(root.as_path()), cap, dry_run);
            for p in &g.removed {
                println!(
                    "{}{}",
                    if dry_run { "would remove " } else { "removed " },
                    p.display()
                );
            }
            for f in &g.failed {
                eprintln!("roundabout: {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!("roundabout: {} 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!("roundabout: {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!("roundabout: {msg}");
            1
        }
        Ok(None) => {
            eprintln!("roundabout: nothing running");
            1
        }
        Ok(Some(_)) => 1,
        Err(e) => {
            eprintln!("roundabout: {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!("roundabout: 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!("roundabout: {e}");
                return 2;
            }
        }
    };
    let mut s = match client::connect_or_spawn() {
        Ok(s) => s,
        Err(e) => {
            eprintln!("roundabout: can't reach the daemon: {e}");
            return 1;
        }
    };
    if let Err(e) = proto::send(&mut s, &Request::Expect { session, need }) {
        eprintln!("roundabout: {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!("roundabout: {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, `roundabout kill {lease}` or `roundabout 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 = "roundabout", 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	/// leave other projects alone
107	#[pound(long)]
108	project: bool,
109	#[pound(short = 'n', long)]
110	dry_run: bool,
111	},
112	/// never gc this
113	Pin {
114	path: PathBuf,
115	},
116	Unpin {
117	path: PathBuf,
118	},
119	}
120
121	fn artifacts(cfg: &Config, action: Artifacts) -> i32 {
122	let cwd = std::env::current_dir().unwrap_or_else(\|_\| "/".into());
123	let root = paths::project_root(&cwd);
124	match action {
125	Artifacts::Check { fresh } => {
126	let r = artifacts::check(cfg, Some(&root), fresh);
127	if let Some((root, size)) = &r.project {
128	println!(
129	"{}: {size} of {}",
130	root.display(),
131	cfg.artifacts.project_cap
132	);
133	}
134	println!("everything: {} of {}", r.global, cfg.artifacts.global_cap);
135	if r.over() {
136	println!("over. `{}`", r.remedy());
137	1
138	} else {
139	0
140	}
141	}
142	Artifacts::List => {
143	let r = artifacts::check(cfg, Some(&root), false);
144	let mut rows: Vec<_> = r.entries.iter().collect();
145	rows.sort_by_key(\|(_, e)\| std::cmp::Reverse(e.size));
146	for (path, e) in rows {
147	println!(
148	"{:>10} {}{}",
149	e.size.to_string(),
150	path.display(),
151	if e.pinned { " (pinned)" } else { "" }
152	);
153	}
154	0
155	}
156	Artifacts::Gc {
157	cap,
158	project,
159	dry_run,
160	} => {
161	let g = artifacts::gc(cfg, project.then_some(root.as_path()), cap, dry_run);
162	for p in &g.removed {
163	println!(
164	"{}{}",
165	if dry_run { "would remove " } else { "removed " },
166	p.display()
167	);
168	}
169	for f in &g.failed {
170	eprintln!("roundabout: {f}");
171	}
172	println!(
173	"{} {}",
174	if dry_run { "would free" } else { "freed" },
175	g.freed
176	);
177	i32::from(!g.failed.is_empty())
178	}
179	Artifacts::Pin { path } \| Artifacts::Unpin { path } if !path.exists() => {
180	eprintln!("roundabout: {} doesn't exist", path.display());
181	1
182	}
183	Artifacts::Pin { path } => report(artifacts::pin(&path, true)),
184	Artifacts::Unpin { path } => report(artifacts::pin(&path, false)),
185	}
186	}
187
188	fn report(r: Result<(), String>) -> i32 {
189	r.map_or_else(
190	\|e\| {
191	eprintln!("roundabout: {e}");
192	1
193	},
194	\|()\| 0,
195	)
196	}
197
198	fn ask(req: &Request) -> io::Result<Option<Event>> {
199	let Some(mut s) = client::connect() else {
200	return Ok(None);
201	};
202	proto::send(&mut s, req)?;
203	Reader::new(&s)?.recv()
204	}
205
206	fn reply(ev: io::Result<Option<Event>>) -> i32 {
207	match ev {
208	Ok(Some(Event::Ok)) => 0,
209	Ok(Some(Event::Err { msg })) => {
210	eprintln!("roundabout: {msg}");
211	1
212	}
213	Ok(None) => {
214	eprintln!("roundabout: nothing running");
215	1
216	}
217	Ok(Some(_)) => 1,
218	Err(e) => {
219	eprintln!("roundabout: {e}");
220	1
221	}
222	}
223	}
224
225	fn agent_session(session: Option<String>) -> Option<String> {
226	session.or_else(\|\| {
227	[
228	"ROUNDABOUT_SESSION_ID",
229	"CLAUDE_CODE_SESSION_ID",
230	"CODEX_THREAD_ID",
231	"CODEX_SESSION_ID",
232	]
233	.into_iter()
234	.find_map(\|key\| std::env::var(key).ok().filter(\|value\| !value.is_empty()))
235	})
236	}
237
238	fn expect(cfg: &Config, session: Option<String>, done: bool, a: &run::Args) -> i32 {
239	let Some(session) = session.filter(\|s\| !s.is_empty()) else {
240	eprintln!("roundabout: no agent session id, use --session to name one");
241	return 2;
242	};
243	let need = if done {
244	None
245	} else {
246	match run::need(cfg, a) {
247	Ok(n) => Some(n),
248	Err(e) => {
249	eprintln!("roundabout: {e}");
250	return 2;
251	}
252	}
253	};
254	let mut s = match client::connect_or_spawn() {
255	Ok(s) => s,
256	Err(e) => {
257	eprintln!("roundabout: can't reach the daemon: {e}");
258	return 1;
259	}
260	};
261	if let Err(e) = proto::send(&mut s, &Request::Expect { session, need }) {
262	eprintln!("roundabout: {e}");
263	return 1;
264	}
265	reply(Reader::new(&s).and_then(\|mut r\| r.recv()))
266	}
267
268	fn status(json: bool) -> i32 {
269	let Ok(Some(Event::Status(st))) = ask(&Request::Status) else {
270	if json {
271	println!("null");
272	} else {
273	println!("idle");
274	}
275	return 0;
276	};
277	if json {
278	println!("{}", serde_json::to_string(&st).unwrap_or_default());
279	return 0;
280	}
281	println!(
282	"ram {} free of {}, cores {} free of {}",
283	st.ram_free, st.ram_budget, st.cores_free, st.cores
284	);
285	for l in &st.leases {
286	let state = if l.running {
287	format!(
288	"running {}{}",
289	Dur::secs(l.elapsed),
290	if l.overdue { " OVERDUE" } else { "" }
291	)
292	} else {
293	format!(
294	"#{} {}",
295	l.position,
296	l.waiting.map(\|r\| r.to_string()).unwrap_or_default()
297	)
298	};
299	let cores = if l.granted > l.need.cores {
300	format!("{}+{}", l.need.cores, l.granted - l.need.cores)
301	} else {
302	l.need.cores.to_string()
303	};
304	println!(
305	"{:>4} {}{} {} ram {} cores {} rss {} {}",
306	l.id,
307	if l.heavy { "H " } else { "" },
308	state,
309	Dur(l.need.time),
310	l.need.ram,
311	cores,
312	l.rss,
313	l.cmd.join(" ")
314	);
315	}
316	for i in &st.intents {
317	println!(
318	" expect {} ram {} cores {}, {} left",
319	&i.session[..i.session.len().min(8)],
320	i.need.ram,
321	i.need.cores,
322	Dur::secs(i.expires_in)
323	);
324	}
325	0
326	}
327
328	fn suggest(cfg: &Config, cmd: &[String]) -> i32 {
329	let cwd = std::env::current_dir().unwrap_or_else(\|_\| "/".into());
330	let key = history::key(&paths::project_root(&cwd), cmd);
331	let Some(rec) = history::lookup(&key) else {
332	println!("never seen it, pick a --size");
333	return 1;
334	};
335	let base = match run::base(cfg) {
336	Ok(b) => b,
337	Err(e) => {
338	eprintln!("roundabout: {e}");
339	return 2;
340	}
341	};
342	let n = history::suggest(&rec, base, cmd);
343	println!(
344	"--ram {} --cores {} --time {} --disk {}",
345	n.ram,
346	n.cores,
347	Dur(n.time),
348	n.disk
349	);
350	0
351	}
352
353	fn watch(lease: u64) -> i32 {
354	let Some(mut s): Option<UnixStream> = client::connect() else {
355	return 0;
356	};
357	if proto::send(&mut s, &Request::Watch { lease }).is_err() {
358	return 0;
359	}
360	if matches!(
361	Reader::new(&s).and_then(\|mut r\| r.recv::<Event>()),
362	Ok(Some(Event::Fired { overdue: true }))
363	) {
364	println!("lease {lease} is overdue, `roundabout kill {lease}` or `roundabout extend {lease} <dur>`");
365	3
366	} else {
367	println!("lease {lease} is done");
368	0
369	}
370	}
371
372	pub fn dispatch(cfg: &Config, cmd: Rabt) -> i32 {
373	match cmd {
374	Rabt::Run {
375	size,
376	ram,
377	disk,
378	cores,
379	time,
380	deadline,
381	max_wait,
382	session,
383	exact,
384	cmd,
385	} => run::run(
386	cfg,
387	&run::Args {
388	size,
389	ram,
390	disk,
391	cores,
392	time,
393	deadline,
394	max_wait,
395	session: agent_session(session),
396	exact,
397	cmd,
398	},
399	),
400	Rabt::Expect {
401	size,
402	ram,
403	cores,
404	done,
405	session,
406	} => expect(
407	cfg,
408	agent_session(session),
409	done,
410	&run::Args {
411	size,
412	ram,
413	cores,
414	..run::Args::default()
415	},
416	),
417	Rabt::Status { json } => status(json),
418	Rabt::Suggest { cmd } => suggest(cfg, &cmd),
419	Rabt::Kill { lease } => reply(ask(&Request::Kill { lease })),
420	Rabt::Extend { lease, by } => reply(ask(&Request::Extend { lease, by })),
421	Rabt::Watch { lease } => watch(lease),
422	Rabt::Artifacts { action } => artifacts(cfg, action),
423	Rabt::Hook { event } => hook::main(cfg, &event),
424	Rabt::Daemon => daemon::main(Config::load().unwrap_or_default()),
425	}
426	}