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