1use std::collections::{BTreeMap, BTreeSet};
2use std::time::{Duration, SystemTime, UNIX_EPOCH};
3
4use anyhow::Result;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7
8use crate::process::pid_alive;
9use crate::providers;
10use crate::status::{
11 AgentSession, AgentStatus, Source, cwd_field, string_field, title_field, u32_field,
12};
13
14pub const PLUGIN_STALE: Duration = Duration::from_secs(5);
15pub const SESSION_RETENTION: Duration = Duration::from_secs(7 * 24 * 60 * 60);
16
17#[derive(Debug, Clone)]
18pub enum Change {
19 Hooks { provider: String },
20 Snapshot { provider: String, instance: String },
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
24#[serde(rename_all = "lowercase")]
25pub enum SessionKind {
26 Hook,
27 Plugin,
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct PluginSnapshot {
32 #[serde(default)]
33 pub status: BTreeMap<String, String>,
34 #[serde(default)]
35 pub blocking: Vec<String>,
36 #[serde(default)]
37 pub titles: BTreeMap<String, String>,
38 #[serde(default)]
39 pub cwd: Option<String>,
40 #[serde(default)]
41 pub cmdline: Vec<String>,
42 #[serde(default)]
43 pub pid: Option<u32>,
44 #[serde(default)]
45 pub last_report_ms: u64,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct Store {
50 #[serde(default)]
51 pub last_heartbeat_ms: u64,
52 #[serde(default)]
53 pub previous_heartbeat_ms: Option<u64>,
54 #[serde(default)]
55 pub hooks: BTreeMap<String, BTreeMap<String, AgentSession>>,
56 #[serde(default)]
57 pub snapshots: BTreeMap<String, BTreeMap<String, PluginSnapshot>>,
58 #[serde(default)]
59 pub removed: BTreeMap<String, BTreeMap<String, u64>>,
60}
61
62impl Default for Store {
63 fn default() -> Self {
64 Self {
65 last_heartbeat_ms: now_ms(),
66 previous_heartbeat_ms: None,
67 hooks: BTreeMap::new(),
68 snapshots: BTreeMap::new(),
69 removed: BTreeMap::new(),
70 }
71 }
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct ListedSession {
76 pub provider: String,
77 pub session_id: String,
78 pub status: AgentStatus,
79 pub source: Source,
80 pub cwd: Option<String>,
81 pub cmdline: Vec<String>,
82 pub pid: Option<u32>,
83 pub last_report_ms: u64,
84 pub kind: SessionKind,
85 #[serde(skip_serializing_if = "Option::is_none")]
86 pub parent_id: Option<String>,
87 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
88 pub exited: bool,
89 #[serde(default, skip_serializing_if = "Option::is_none")]
90 pub title: Option<String>,
91}
92
93impl Store {
94 pub fn on_server_start(&mut self) {
95 let now = now_ms();
96 if self.last_heartbeat_ms > 0 {
97 self.previous_heartbeat_ms = Some(self.last_heartbeat_ms);
98 }
99 self.last_heartbeat_ms = now;
100 }
101
102 pub fn heartbeat(&mut self) {
103 self.last_heartbeat_ms = now_ms();
104 }
105
106 pub fn update(&mut self, provider: &str, payload: Value) -> Result<Change> {
107 if !payload.is_object() {
108 anyhow::bail!("provide a JSON object");
109 }
110 if payload.get("hook_event_name").is_some() || payload.get("hookEventName").is_some() {
111 let sid = crate::status::session_key(&payload);
112 if let Some(sid) = &sid {
113 self.clear_removed(provider, sid);
114 }
115 let bucket = self.hooks.entry(provider.to_string()).or_default();
116 providers::apply_hook(provider, bucket, &payload);
117 if let Some(sid) = sid {
118 if let Some(session) = bucket.get_mut(&sid) {
119 touch_session(session, &payload, provider);
120 if session.exited {
121 bucket.remove(&sid);
122 }
123 } else if crate::status::string_field(
124 &payload,
125 &["hook_event_name", "hookEventName"],
126 )
127 .is_some_and(|name| name.eq_ignore_ascii_case("SessionEnd"))
128 {
129 }
131 }
132 return Ok(Change::Hooks {
133 provider: provider.to_string(),
134 });
135 }
136 let instance = payload
137 .get("id")
138 .and_then(Value::as_str)
139 .filter(|s| !s.is_empty())
140 .ok_or_else(|| anyhow::anyhow!("provide a string id"))?;
141 let raw_status = payload
142 .get("status")
143 .ok_or_else(|| anyhow::anyhow!("status must be an object"))?;
144 let obj = raw_status
145 .as_object()
146 .ok_or_else(|| anyhow::anyhow!("status must be an object"))?;
147 let mut status = BTreeMap::new();
148 for (sid, kind) in obj {
149 let kind = kind
150 .as_str()
151 .ok_or_else(|| anyhow::anyhow!("status values must be idle, busy, or retry"))?;
152 if !matches!(kind, "idle" | "busy" | "retry") {
153 anyhow::bail!("status values must be idle, busy, or retry");
154 }
155 status.insert(sid.clone(), kind.to_string());
156 }
157 for sid in status.keys() {
158 self.clear_removed(provider, sid);
159 }
160 let blocking = match payload.get("blocking") {
161 None => Vec::new(),
162 Some(Value::Array(items)) => {
163 let mut out = Vec::new();
164 for item in items {
165 let sid = item
166 .as_str()
167 .ok_or_else(|| anyhow::anyhow!("blocking must be an array of strings"))?;
168 out.push(sid.to_string());
169 }
170 out
171 }
172 Some(_) => anyhow::bail!("blocking must be an array of strings"),
173 };
174 let titles = match payload.get("titles") {
175 None => BTreeMap::new(),
176 Some(Value::Object(items)) => {
177 let mut out = BTreeMap::new();
178 for (sid, title) in items {
179 let title = title
180 .as_str()
181 .ok_or_else(|| anyhow::anyhow!("titles must be an object of strings"))?;
182 out.insert(sid.clone(), title.to_string());
183 }
184 out
185 }
186 Some(_) => anyhow::bail!("titles must be an object of strings"),
187 };
188 let pid = u32_field(&payload, &["pid"]).or_else(|| instance.parse().ok());
189 let snapshot = PluginSnapshot {
190 status,
191 blocking,
192 titles,
193 cwd: cwd_field(&payload),
194 cmdline: string_list(&payload, "cmdline"),
195 pid,
196 last_report_ms: now_ms(),
197 };
198 self.snapshots
199 .entry(provider.to_string())
200 .or_default()
201 .insert(instance.to_string(), snapshot);
202 Ok(Change::Snapshot {
203 provider: provider.to_string(),
204 instance: instance.to_string(),
205 })
206 }
207
208 pub fn mark_removed(&mut self, provider: &str, session_id: &str) -> u64 {
209 let at = now_ms();
210 self.removed
211 .entry(provider.to_string())
212 .or_default()
213 .insert(session_id.to_string(), at);
214 at
215 }
216
217 pub fn clear_removed(&mut self, provider: &str, session_id: &str) {
218 if let Some(bucket) = self.removed.get_mut(provider) {
219 bucket.remove(session_id);
220 }
221 }
222
223 pub fn is_removed(&self, provider: &str, session_id: &str) -> bool {
224 self.removed
225 .get(provider)
226 .is_some_and(|bucket| bucket.contains_key(session_id))
227 }
228
229 pub fn discover(&mut self, ctx: &crate::paths::Context) -> bool {
230 let claude =
231 crate::providers::discover_claude(ctx, self.hooks.entry("claude".into()).or_default());
232 let codex = crate::providers::drop_archived_codex(
233 ctx,
234 self.hooks.entry("codex".into()).or_default(),
235 );
236 claude || codex
237 }
238
239 pub fn listed(&self) -> Vec<ListedSession> {
240 let mut out = Vec::new();
241 for (provider, sessions) in &self.hooks {
242 for (sid, session) in sessions {
243 let mut cmdline = session.cmdline.clone();
244 if cmdline.is_empty() {
245 cmdline = providers::resume_cmd(provider, sid);
246 }
247 out.push(ListedSession {
248 provider: provider.clone(),
249 session_id: sid.clone(),
250 status: session.status,
251 source: session.source,
252 cwd: session.cwd.clone(),
253 cmdline,
254 pid: session.pid,
255 last_report_ms: session.last_report_ms,
256 kind: SessionKind::Hook,
257 parent_id: session.parent_id.clone(),
258 exited: session.exited,
259 title: session.title.clone(),
260 });
261 }
262 }
263 for (provider, instances) in &self.snapshots {
264 for snapshot in instances.values() {
265 let blocking: std::collections::BTreeSet<_> =
266 snapshot.blocking.iter().cloned().collect();
267 for (sid, kind) in &snapshot.status {
268 let status = if blocking.contains(sid) {
269 AgentStatus::Waiting
270 } else if matches!(kind.as_str(), "busy" | "retry") {
271 AgentStatus::Working
272 } else {
273 AgentStatus::Idle
274 };
275 let mut cmdline = snapshot.cmdline.clone();
276 if cmdline.is_empty() {
277 cmdline = providers::resume_cmd(provider, sid);
278 }
279 out.push(ListedSession {
280 provider: provider.clone(),
281 session_id: sid.clone(),
282 status,
283 source: Source::Cli,
284 cwd: snapshot.cwd.clone(),
285 cmdline,
286 pid: snapshot.pid,
287 last_report_ms: snapshot.last_report_ms,
288 kind: SessionKind::Plugin,
289 parent_id: None,
290 exited: false,
291 title: snapshot.titles.get(sid).cloned(),
292 });
293 }
294 }
295 }
296 out.retain(|session| !self.is_removed(&session.provider, &session.session_id));
297 let mut best: BTreeMap<(String, String, SessionKind), ListedSession> = BTreeMap::new();
298 for session in out {
299 let key = (
300 session.provider.clone(),
301 session.session_id.clone(),
302 session.kind,
303 );
304 match best.get(&key) {
305 Some(existing) if existing.last_report_ms >= session.last_report_ms => {}
306 _ => {
307 best.insert(key, session);
308 }
309 }
310 }
311 let mut out: Vec<ListedSession> = best.into_values().collect();
312 out.sort_by(|a, b| (&a.provider, &a.session_id).cmp(&(&b.provider, &b.session_id)));
313 out
314 }
315
316 pub fn active(&self) -> Vec<ListedSession> {
317 let now = now_ms();
318 self.listed()
319 .into_iter()
320 .filter(|session| session.is_active(now))
321 .collect()
322 }
323
324 pub fn resumable(&self, idle: Option<Duration>) -> Vec<ListedSession> {
325 let now = now_ms();
326 self.listed()
327 .into_iter()
328 .filter(|session| session.is_resumable(self, now, idle))
329 .collect()
330 }
331
332 pub fn prune(&mut self, now_ms: u64, retention: Duration) -> bool {
333 let cutoff = now_ms.saturating_sub(retention.as_millis() as u64);
334 let mut changed = false;
335
336 for bucket in self.hooks.values_mut() {
337 let before = bucket.len();
338 bucket.retain(|_, session| {
339 session.pid.is_some_and(pid_alive) || session.last_report_ms >= cutoff
340 });
341 changed |= bucket.len() != before;
342 }
343 self.hooks.retain(|_, bucket| !bucket.is_empty());
344
345 for bucket in self.snapshots.values_mut() {
346 let before = bucket.len();
347 bucket.retain(|_, snapshot| {
348 snapshot.pid.is_some_and(pid_alive) || snapshot.last_report_ms >= cutoff
349 });
350 changed |= bucket.len() != before;
351 }
352 self.snapshots.retain(|_, bucket| !bucket.is_empty());
353
354 let mut live: BTreeSet<(String, String)> = BTreeSet::new();
355 for (provider, bucket) in &self.hooks {
356 for sid in bucket.keys() {
357 live.insert((provider.clone(), sid.clone()));
358 }
359 }
360 for (provider, bucket) in &self.snapshots {
361 for snapshot in bucket.values() {
362 for sid in snapshot.status.keys() {
363 live.insert((provider.clone(), sid.clone()));
364 }
365 }
366 }
367 for (provider, bucket) in self.removed.iter_mut() {
368 let before = bucket.len();
369 bucket.retain(|sid, _| live.contains(&(provider.clone(), sid.clone())));
370 changed |= bucket.len() != before;
371 }
372 self.removed.retain(|_, bucket| !bucket.is_empty());
373 changed
374 }
375}
376
377impl ListedSession {
378 pub fn is_active(&self, now_ms: u64) -> bool {
379 if self.exited || self.parent_id.is_some() {
380 return false;
381 }
382 match self.kind {
383 SessionKind::Plugin => {
384 now_ms.saturating_sub(self.last_report_ms) <= PLUGIN_STALE.as_millis() as u64
385 }
386 SessionKind::Hook => {
387 self.status.is_busy()
388 || self.pid.is_some_and(pid_alive)
389 || (self.source == Source::Desktop && !self.exited)
390 }
391 }
392 }
393
394 pub fn process_attached(&self) -> bool {
395 self.pid.is_some_and(pid_alive)
396 }
397
398 pub fn is_resumable(&self, store: &Store, now_ms: u64, idle: Option<Duration>) -> bool {
399 if self.exited || self.parent_id.is_some() {
400 return false;
401 }
402 if self.process_attached() {
403 return false;
404 }
405 if self.cwd.as_ref().is_none_or(|cwd| cwd.is_empty()) {
406 return false;
407 }
408 if self.status.is_busy() {
409 return true;
410 }
411 if !matches!(self.status, AgentStatus::Idle | AgentStatus::Done) {
412 return false;
413 }
414 let Some(idle) = idle else {
415 return false;
416 };
417 idle_age_ms(self, store, now_ms) <= idle.as_millis() as u64
418 }
419}
420
421fn idle_age_ms(session: &ListedSession, store: &Store, now_ms: u64) -> u64 {
422 let anchor = store.previous_heartbeat_ms.unwrap_or(now_ms);
423 if session.last_report_ms > anchor {
424 now_ms.saturating_sub(session.last_report_ms)
425 } else {
426 anchor.saturating_sub(session.last_report_ms)
427 }
428}
429
430fn touch_session(session: &mut AgentSession, payload: &Value, provider: &str) {
431 session.last_report_ms = now_ms();
432 if let Some(cwd) = cwd_field(payload) {
433 session.cwd = Some(cwd);
434 }
435 if let Some(title) = title_field(payload) {
436 session.title = Some(title);
437 }
438 if let Some(pid) = u32_field(payload, &["pid"]) {
439 session.pid = Some(pid);
440 }
441 if session.cmdline.is_empty()
442 && let Some(sid) = crate::status::session_key(payload)
443 {
444 session.cmdline = providers::resume_cmd(provider, &sid);
445 }
446 let name = string_field(payload, &["hook_event_name", "hookEventName"]).unwrap_or("");
447 if name.eq_ignore_ascii_case("SessionEnd") {
448 session.exited = true;
449 }
450}
451
452fn string_list(payload: &Value, key: &str) -> Vec<String> {
453 payload
454 .get(key)
455 .and_then(Value::as_array)
456 .map(|items| {
457 items
458 .iter()
459 .filter_map(Value::as_str)
460 .map(str::to_string)
461 .collect()
462 })
463 .unwrap_or_default()
464}
465
466pub fn now_ms() -> u64 {
467 SystemTime::now()
468 .duration_since(UNIX_EPOCH)
469 .map(|d| d.as_millis() as u64)
470 .unwrap_or(0)
471}
472
473#[cfg(test)]
474#[path = "store_tests.rs"]
475mod tests;