1use serde::Serialize;
7use std::collections::BTreeMap;
8use std::path::{Path, PathBuf};
9use std::sync::{Mutex, OnceLock};
10
11#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
12pub struct SessionGit {
13 pub workspace_id: String,
14 pub instance_id: u64,
15 pub provider: String,
16 pub branch: String,
17 pub head: String,
18 pub dirty: u32,
20 pub ahead: u32,
22 pub root: String,
23}
24
25#[derive(Clone, Debug)]
26struct Live {
27 pid: u32,
28 provider: String,
29 dir: Option<PathBuf>,
32}
33
34#[derive(Clone, Debug, PartialEq, Eq)]
36pub struct Lifecycle {
37 pub workspace: String,
38 pub instance: u64,
39 pub state: &'static str,
41 pub dir: Option<PathBuf>,
43 pub public_key: Option<String>,
46}
47
48static LIFECYCLE: OnceLock<tokio::sync::mpsc::UnboundedSender<Lifecycle>> = OnceLock::new();
49
50pub fn set_lifecycle_sink(tx: tokio::sync::mpsc::UnboundedSender<Lifecycle>) {
52 let _ = LIFECYCLE.set(tx);
53}
54
55fn emit(l: Lifecycle) {
56 if let Some(tx) = LIFECYCLE.get() {
57 let _ = tx.send(l);
58 }
59}
60
61static LIVE: OnceLock<Mutex<BTreeMap<(String, u64), Live>>> = OnceLock::new();
62
63fn live() -> std::sync::MutexGuard<'static, BTreeMap<(String, u64), Live>> {
64 LIVE.get_or_init(Default::default).lock().unwrap_or_else(|p| p.into_inner())
65}
66
67pub fn running(workspace: &str, instance: u64, provider: &str, pid: u32) {
69 live().insert((workspace.to_owned(), instance), Live { pid, provider: provider.to_owned(), dir: cwd_of(pid) });
70 emit(Lifecycle { workspace: workspace.to_owned(), instance, state: "live", dir: None, public_key: None });
71}
72
73pub fn provider_hint(workspace: &str, instance: u64, provider: &str) {
75 if let Some(l) = live().get_mut(&(workspace.to_owned(), instance)) {
76 if l.provider != provider {
77 l.provider = provider.to_owned();
78 }
79 }
80}
81
82pub fn exited(workspace: &str, instance: u64) {
84 if let Some(l) = live().remove(&(workspace.to_owned(), instance)) {
85 emit(Lifecycle { workspace: workspace.to_owned(), instance, state: "gone", dir: l.dir, public_key: None });
86 }
87}
88
89pub fn session_key(proxy_session: &str, public_key_hex: &str) {
92 let Some(instance) = proxy_session.rsplit('-').next().and_then(|i| i.parse::<u64>().ok()) else { return };
93 if public_key_hex.is_empty() {
94 return;
95 }
96 let workspace = live().keys().find(|(_, i)| *i == instance).map(|(w, _)| w.clone());
97 if let Some(workspace) = workspace {
98 emit(Lifecycle { workspace, instance, state: "live", dir: None, public_key: Some(public_key_hex.to_owned()) });
99 }
100}
101
102pub fn read_session_repo(dir: &Path) -> Option<serde_json::Value> {
105 let (branch, head, dirty, ahead, root) = read_repo_full(dir).or_else(|| task_repo_below(dir))?;
106 Some(serde_json::json!({ "branch": branch, "head": head, "dirty": dirty, "ahead": ahead, "root": root }))
107}
108
109fn cwd_of(pid: u32) -> Option<PathBuf> {
110 if cfg!(unix) {
111 std::fs::read_link(format!("/proc/{pid}/cwd")).ok()
112 } else {
113 None
114 }
115}
116
117fn git(dir: &Path, args: &[&str]) -> Option<String> {
118 let mut cmd = std::process::Command::new("git");
119 cmd.arg("-C").arg(dir).args(args).stdin(std::process::Stdio::null()).stderr(std::process::Stdio::null());
120 crate::windows_process::hide_console(&mut cmd);
121 let out = cmd.output().ok()?;
122 out.status.success().then(|| String::from_utf8_lossy(&out.stdout).into_owned())
123}
124
125pub fn read_repo(dir: &Path) -> Option<(String, String, u32, String)> {
127 read_repo_full(dir).map(|(b, h, d, _, r)| (b, h, d, r))
128}
129
130pub fn read_repo_full(dir: &Path) -> Option<(String, String, u32, u32, String)> {
132 let (branch, head, dirty, root) = read_repo_basic(dir)?;
133 let base = ["main", "master"].into_iter().find(|b| git(dir, &["rev-parse", "--verify", "--quiet", &format!("refs/heads/{b}")]).is_some());
134 let ahead = base
135 .filter(|b| *b != branch)
136 .and_then(|b| git(dir, &["rev-list", "--count", "HEAD", &format!("^{b}")]))
137 .and_then(|s| s.trim().parse().ok())
138 .unwrap_or(0);
139 Some((branch, head, dirty, ahead, root))
140}
141
142fn read_repo_basic(dir: &Path) -> Option<(String, String, u32, String)> {
143 let root = git(dir, &["rev-parse", "--show-toplevel"])?.trim().to_owned();
144 let branch = git(dir, &["rev-parse", "--abbrev-ref", "HEAD"]).map(|s| s.trim().to_owned()).unwrap_or_default();
145 let head = git(dir, &["rev-parse", "--short=12", "HEAD"]).map(|s| s.trim().to_owned()).unwrap_or_default();
146 let dirty = git(dir, &["status", "--porcelain", "--untracked-files=normal"]).map(|s| s.lines().take(10_000).count() as u32).unwrap_or(0);
147 Some((branch, head, dirty, root))
148}
149
150fn task_repo_below(dir: &Path) -> Option<(String, String, u32, u32, String)> {
153 let mut entries: Vec<_> = std::fs::read_dir(dir).ok()?.flatten().take(256).map(|e| e.path()).filter(|p| p.join(".git").exists()).collect();
154 entries.sort();
155 entries.into_iter().find_map(|p| read_repo_full(&p).filter(|r| r.0.starts_with("task-")))
156}
157
158pub fn snapshot() -> Vec<SessionGit> {
160 let sessions: Vec<_> = live().iter().map(|(k, v)| (k.clone(), v.clone())).collect();
161 let out: Vec<SessionGit> = sessions
162 .into_iter()
163 .filter_map(|((ws, inst), l)| {
164 let dir = cwd_of(l.pid)?;
165 let (branch, head, dirty, ahead, root) = read_repo_full(&dir).or_else(|| task_repo_below(&dir))?;
166 Some(SessionGit { workspace_id: ws, instance_id: inst, provider: l.provider, branch, head, dirty, ahead, root })
167 })
168 .collect();
169 let mut live = live();
170 for s in &out {
171 if let Some(l) = live.get_mut(&(s.workspace_id.clone(), s.instance_id)) {
172 l.dir = Some(PathBuf::from(&s.root));
173 }
174 }
175 out
176}
177
178#[cfg(test)]
179mod tests {
180 use super::*;
181
182 #[test]
183 fn reads_branch_and_dirty() {
184 let d = std::env::temp_dir().join(format!("gw-{}", std::process::id()));
185 let _ = std::fs::remove_dir_all(&d);
186 std::fs::create_dir_all(&d).unwrap();
187 let run = |a: &[&str]| assert!(std::process::Command::new("git").arg("-C").arg(&d).args(a).output().unwrap().status.success());
188 run(&["init", "-q", "-b", "feat-x"]);
189 run(&["-c", "core.hooksPath=/dev/null", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "--allow-empty", "-m", "i"]);
190 std::fs::write(d.join("a.txt"), "x").unwrap();
191 let (branch, head, dirty, _) = read_repo(&d).unwrap();
192 run(&["branch", "main"]);
193 run(&["-c", "core.hooksPath=/dev/null", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "--allow-empty", "-m", "j"]);
194 assert_eq!(read_repo_full(&d).unwrap().3, 1);
195 assert_eq!(branch, "feat-x");
196 assert_eq!(head.len(), 12);
197 assert_eq!(dirty, 1);
198 let _ = std::fs::remove_dir_all(d);
199 }
200
201 #[test]
202 fn finds_task_repo_below_workspace() {
203 let ws = std::env::temp_dir().join(format!("gw-ws-{}", std::process::id()));
204 let _ = std::fs::remove_dir_all(&ws);
205 let r = ws.join("repo");
206 std::fs::create_dir_all(&r).unwrap();
207 let run = |a: &[&str]| assert!(std::process::Command::new("git").arg("-C").arg(&r).args(a).output().unwrap().status.success());
208 run(&["init", "-q", "-b", "main"]);
209 run(&["-c", "core.hooksPath=/dev/null", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "--allow-empty", "-m", "i"]);
210 assert!(task_repo_below(&ws).is_none());
211 run(&["checkout", "-q", "-b", "task-9"]);
212 run(&["-c", "core.hooksPath=/dev/null", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "--allow-empty", "-m", "w"]);
213 let f = task_repo_below(&ws).unwrap();
214 assert_eq!((f.0.as_str(), f.2, f.3), ("task-9", 0, 1));
215 let _ = std::fs::remove_dir_all(ws);
216 }
217
218 #[test]
219 fn lifecycle_is_emitted_with_the_last_dir() {
220 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
221 set_lifecycle_sink(tx);
222 running("lc", 91, "codex", std::process::id());
223 session_key("inc-1-91", "ab");
224 exited("lc", 91);
225 exited("lc", 91);
226 let mut got = Vec::new();
227 while let Ok(l) = rx.try_recv() {
228 if l.workspace == "lc" {
229 got.push(l);
230 }
231 }
232 assert_eq!(got.iter().map(|l| l.state).collect::<Vec<_>>(), ["live", "live", "gone"], "a second exit is not a second gone");
233 assert_eq!(got[1].public_key.as_deref(), Some("ab"));
234 if cfg!(target_os = "linux") {
235 assert_eq!(got[2].dir, std::env::current_dir().ok());
236 }
237 let here = std::env::current_dir().unwrap();
238 assert_eq!(read_session_repo(&here).is_some(), read_repo(&here).is_some());
239 }
240
241 #[test]
242 fn live_session_snapshot() {
243 running("w", 77, "claude", std::process::id());
244 let here = std::env::current_dir().unwrap();
245 let in_repo = read_repo(&here).is_some();
246 assert_eq!(snapshot().iter().any(|s| s.instance_id == 77), in_repo && cfg!(target_os = "linux"));
247 exited("w", 77);
248 assert!(!snapshot().iter().any(|s| s.instance_id == 77));
249 }
250}