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