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

agent_berth/
status.rs
