Skip to main content

agent_berth/
status.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
4#[serde(rename_all = "kebab-case")]
5pub enum AgentEventKind {
6    SessionStart,
7    PromptSubmit,
8    ToolStart,
9    ToolComplete,
10    PermissionRequest,
11    QuestionAsked,
12    Notification,
13    Stop,
14    SessionEnd,
15}
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
18#[serde(rename_all = "lowercase")]
19pub enum AgentStatus {
20    #[default]
21    Idle,
22    Working,
23    Waiting,
24    Done,
25}
26
27impl AgentStatus {
28    pub fn as_str(self) -> &'static str {
29        match self {
30            Self::Idle => "idle",
31            Self::Working => "working",
32            Self::Waiting => "waiting",
33            Self::Done => "done",
34        }
35    }
36
37    pub fn is_busy(self) -> bool {
38        matches!(self, Self::Working | Self::Waiting)
39    }
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
43#[serde(rename_all = "lowercase")]
44pub enum Source {
45    #[default]
46    Cli,
47    Desktop,
48}
49
50impl Source {
51    pub fn as_str(self) -> &'static str {
52        match self {
53            Self::Cli => "cli",
54            Self::Desktop => "desktop",
55        }
56    }
57
58    pub fn from_label(value: &str) -> Self {
59        if value.to_ascii_lowercase().contains("desktop") {
60            Self::Desktop
61        } else {
62            Self::Cli
63        }
64    }
65}
66
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub struct AgentEvent {
69    pub session_id: String,
70    pub kind: AgentEventKind,
71    pub source: Source,
72    pub pid: Option<u32>,
73    pub parent_id: Option<String>,
74    pub background_running: bool,
75}
76
77impl AgentEvent {
78    pub fn new(session_id: impl Into<String>, kind: AgentEventKind) -> Self {
79        Self {
80            session_id: session_id.into(),
81            kind,
82            source: Source::Cli,
83            pid: None,
84            parent_id: None,
85            background_running: false,
86        }
87    }
88}
89
90#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
91pub struct AgentSession {
92    #[serde(default)]
93    pub status: AgentStatus,
94    #[serde(default)]
95    pub source: Source,
96    #[serde(default)]
97    pub pid: Option<u32>,
98    #[serde(default)]
99    pub parent_id: Option<String>,
100    #[serde(default)]
101    pub background_only: bool,
102    #[serde(default)]
103    pub cwd: Option<String>,
104    #[serde(default)]
105    pub cmdline: Vec<String>,
106    #[serde(default)]
107    pub last_report_ms: u64,
108    #[serde(default)]
109    pub exited: bool,
110    #[serde(default)]
111    pub discovered: bool,
112    #[serde(default)]
113    pub title: Option<String>,
114}
115
116impl Default for AgentSession {
117    fn default() -> Self {
118        Self {
119            status: AgentStatus::Idle,
120            source: Source::Cli,
121            pid: None,
122            parent_id: None,
123            background_only: false,
124            cwd: None,
125            cmdline: Vec::new(),
126            last_report_ms: 0,
127            exited: false,
128            discovered: false,
129            title: None,
130        }
131    }
132}
133
134impl AgentSession {
135    pub fn apply(&mut self, event: &AgentEvent) {
136        self.source = event.source;
137        if event.pid.is_some() {
138            self.pid = event.pid;
139        }
140        if event.parent_id.is_some() {
141            self.parent_id = event.parent_id.clone();
142        }
143        match event.kind {
144            AgentEventKind::PromptSubmit => {
145                self.status = AgentStatus::Working;
146                self.background_only = false;
147            }
148            AgentEventKind::ToolStart => {
149                if self.status == AgentStatus::Idle {
150                    self.status = AgentStatus::Working;
151                }
152                if self.status == AgentStatus::Working {
153                    self.background_only = false;
154                }
155            }
156            AgentEventKind::ToolComplete => {
157                if self.status == AgentStatus::Waiting {
158                    self.status = AgentStatus::Working;
159                }
160                if self.status == AgentStatus::Working {
161                    self.background_only = false;
162                }
163            }
164            AgentEventKind::PermissionRequest | AgentEventKind::QuestionAsked => {
165                self.status = AgentStatus::Waiting;
166            }
167            AgentEventKind::Notification => {
168                if self.status == AgentStatus::Working {
169                    self.status = AgentStatus::Waiting;
170                }
171            }
172            AgentEventKind::Stop => {
173                self.status = if event.background_running {
174                    AgentStatus::Working
175                } else {
176                    AgentStatus::Done
177                };
178                self.background_only = event.background_running;
179            }
180            AgentEventKind::SessionStart | AgentEventKind::SessionEnd => {}
181        }
182    }
183}
184
185pub fn prune(
186    sessions: &mut std::collections::BTreeMap<String, AgentSession>,
187    keep: Option<&str>,
188    source: Option<Source>,
189    pid: Option<u32>,
190) {
191    sessions.retain(|sid, item| {
192        if Some(sid.as_str()) == keep || item.status.is_busy() {
193            return true;
194        }
195        if source.is_some() && Some(item.source) != source {
196            return true;
197        }
198        item.pid != pid
199    });
200}
201
202pub fn apply_event(
203    sessions: &mut std::collections::BTreeMap<String, AgentSession>,
204    event: AgentEvent,
205) {
206    let sid = event.session_id.clone();
207    let kind = event.kind;
208    let existing_pid = sessions.get(&sid).and_then(|item| item.pid);
209    let pid = event.pid.or(existing_pid);
210    if kind == AgentEventKind::SessionEnd {
211        sessions.remove(&sid);
212        sessions.retain(|_, child| child.parent_id.as_deref() != Some(sid.as_str()));
213        return;
214    }
215    if kind == AgentEventKind::SessionStart {
216        prune(sessions, Some(&sid), Some(event.source), pid);
217        return;
218    }
219    if !sessions.contains_key(&sid) {
220        if matches!(
221            kind,
222            AgentEventKind::ToolComplete | AgentEventKind::Notification
223        ) {
224            return;
225        }
226        if kind == AgentEventKind::Stop && !event.background_running {
227            return;
228        }
229        sessions.insert(sid.clone(), AgentSession::default());
230    }
231    let parent_id;
232    let source;
233    let session_pid;
234    {
235        let item = sessions.get_mut(&sid).expect("session inserted");
236        item.apply(&event);
237        parent_id = item.parent_id.clone();
238        source = item.source;
239        session_pid = item.pid;
240    }
241    if parent_id.is_some() {
242        if sessions
243            .get(&sid)
244            .is_some_and(|item| item.status == AgentStatus::Done)
245        {
246            sessions.remove(&sid);
247        }
248    } else if matches!(kind, AgentEventKind::PromptSubmit | AgentEventKind::Stop) {
249        prune(sessions, Some(&sid), Some(source), session_pid);
250    }
251}
252
253pub fn session_key(payload: &serde_json::Value) -> Option<String> {
254    let sid = string_field(payload, &["session_id", "sessionId"]);
255    let agent = string_field(payload, &["agent_id", "agentId"]);
256    match (sid, agent) {
257        (Some(sid), Some(agent)) => Some(format!("{sid}:{agent}")),
258        (None, Some(agent)) => Some(agent.to_string()),
259        (Some(sid), None) => Some(sid.to_string()),
260        (None, None) => None,
261    }
262}
263
264pub fn string_field<'a>(payload: &'a serde_json::Value, names: &[&str]) -> Option<&'a str> {
265    let obj = payload.as_object()?;
266    for name in names {
267        match obj.get(*name) {
268            Some(serde_json::Value::String(value)) if !value.is_empty() => return Some(value),
269            _ => {}
270        }
271    }
272    None
273}
274
275pub fn u32_field(payload: &serde_json::Value, names: &[&str]) -> Option<u32> {
276    let obj = payload.as_object()?;
277    for name in names {
278        match obj.get(*name) {
279            Some(serde_json::Value::Number(n)) => {
280                if let Some(v) = n.as_u64() {
281                    return u32::try_from(v).ok();
282                }
283            }
284            Some(serde_json::Value::String(s)) => {
285                if let Ok(v) = s.parse() {
286                    return Some(v);
287                }
288            }
289            _ => {}
290        }
291    }
292    None
293}
294
295pub fn title_field(payload: &serde_json::Value) -> Option<String> {
296    string_field(
297        payload,
298        &[
299            "title",
300            "name",
301            "thread_name",
302            "session_name",
303            "conversation_title",
304            "display_name",
305        ],
306    )
307    .map(str::to_string)
308}
309
310pub fn cwd_field(payload: &serde_json::Value) -> Option<String> {
311    string_field(
312        payload,
313        &[
314            "cwd",
315            "workdir",
316            "workDir",
317            "working_directory",
318            "work_dir",
319            "directory",
320        ],
321    )
322    .map(str::to_string)
323}
324
325#[cfg(test)]
326#[path = "status_tests.rs"]
327mod tests;