| 1 |
|
| 2 |
|
| 3 |
|
| 4 |
|
| 5 |
use crate::event::NewEvent; |
| 6 |
use std::collections::HashMap; |
| 7 |
|
| 8 |
|
| 9 |
pub const ONLINE_MS: u64 = 60_000; |
| 10 |
|
| 11 |
pub const IDLE_AFTER_MS: u64 = 10 * 60_000; |
| 12 |
|
| 13 |
#[derive(Clone, Debug, Default)] |
| 14 |
struct Activity { |
| 15 |
last_request_ms: Option<u64>, |
| 16 |
waiting: u32, |
| 17 |
shims: u32, |
| 18 |
status_ms: Option<u64>, |
| 19 |
blocked: bool, |
| 20 |
paused: bool, |
| 21 |
session: Option<String>, |
| 22 |
state: Option<&'static str>, |
| 23 |
} |
| 24 |
|
| 25 |
|
| 26 |
#[derive(Debug, Default)] |
| 27 |
pub struct Presence { |
| 28 |
agents: HashMap<String, Activity>, |
| 29 |
} |
| 30 |
|
| 31 |
impl Presence { |
| 32 |
|
| 33 |
|
| 34 |
pub fn request(&mut self, agent: &str, now: u64, new_session: bool) -> (String, Option<NewEvent>) { |
| 35 |
let a = self.agents.entry(agent.to_string()).or_default(); |
| 36 |
let lapsed = a.last_request_ms.is_none_or(|t| now.saturating_sub(t) >= IDLE_AFTER_MS); |
| 37 |
a.last_request_ms = Some(now); |
| 38 |
if a.session.is_some() && !lapsed && !new_session { |
| 39 |
return (a.session.clone().unwrap_or_default(), None); |
| 40 |
} |
| 41 |
let id = crate::ids::ulid(); |
| 42 |
a.session = Some(id.clone()); |
| 43 |
let e = NewEvent::system("agent.session_started", serde_json::json!({ "agent": agent, "session": id })) |
| 44 |
.subject(agent) |
| 45 |
.session(id.clone()); |
| 46 |
(id, Some(e)) |
| 47 |
} |
| 48 |
pub fn wait_started(&mut self, agent: &str) { |
| 49 |
self.agents.entry(agent.to_string()).or_default().waiting += 1; |
| 50 |
} |
| 51 |
pub fn wait_ended(&mut self, agent: &str) { |
| 52 |
let a = self.agents.entry(agent.to_string()).or_default(); |
| 53 |
a.waiting = a.waiting.saturating_sub(1); |
| 54 |
} |
| 55 |
pub fn shim_connected(&mut self, agent: &str) { |
| 56 |
self.agents.entry(agent.to_string()).or_default().shims += 1; |
| 57 |
} |
| 58 |
pub fn shim_gone(&mut self, agent: &str) { |
| 59 |
let a = self.agents.entry(agent.to_string()).or_default(); |
| 60 |
a.shims = a.shims.saturating_sub(1); |
| 61 |
} |
| 62 |
|
| 63 |
pub fn status(&mut self, agent: &str, at: Option<u64>) { |
| 64 |
self.agents.entry(agent.to_string()).or_default().status_ms = at; |
| 65 |
} |
| 66 |
pub fn blocked(&mut self, agent: &str, blocked: bool) { |
| 67 |
self.agents.entry(agent.to_string()).or_default().blocked = blocked; |
| 68 |
} |
| 69 |
pub fn paused(&mut self, agent: &str, paused: bool) { |
| 70 |
self.agents.entry(agent.to_string()).or_default().paused = paused; |
| 71 |
} |
| 72 |
|
| 73 |
|
| 74 |
pub fn state(&self, agent: &str, now: u64) -> &'static str { |
| 75 |
self.agents.get(agent).map(|a| compute(a, now)).unwrap_or("offline") |
| 76 |
} |
| 77 |
|
| 78 |
|
| 79 |
|
| 80 |
pub fn tick(&mut self, now: u64) -> Vec<NewEvent> { |
| 81 |
let mut out = Vec::new(); |
| 82 |
let mut names: Vec<&String> = self.agents.keys().collect(); |
| 83 |
names.sort(); |
| 84 |
let names: Vec<String> = names.into_iter().cloned().collect(); |
| 85 |
for name in names { |
| 86 |
let a = self.agents.get_mut(&name).expect("listed"); |
| 87 |
let s = compute(a, now); |
| 88 |
if a.state != Some(s) { |
| 89 |
let from = a.state.unwrap_or("offline"); |
| 90 |
a.state = Some(s); |
| 91 |
if from != s { |
| 92 |
out.push( |
| 93 |
NewEvent::system("agent.presence", serde_json::json!({ "agent": name, "state": s, "from": from })) |
| 94 |
.subject(name.clone()), |
| 95 |
); |
| 96 |
} |
| 97 |
} |
| 98 |
let quiet = a.last_request_ms.is_none_or(|t| now.saturating_sub(t) >= IDLE_AFTER_MS); |
| 99 |
if quiet && a.waiting == 0 { |
| 100 |
if let Some(id) = a.session.take() { |
| 101 |
out.push( |
| 102 |
NewEvent::system("agent.session_ended", serde_json::json!({ "agent": name, "session": id })) |
| 103 |
.subject(name.clone()) |
| 104 |
.session(id), |
| 105 |
); |
| 106 |
} |
| 107 |
} |
| 108 |
} |
| 109 |
out |
| 110 |
} |
| 111 |
} |
| 112 |
|
| 113 |
|
| 114 |
fn compute(a: &Activity, now: u64) -> &'static str { |
| 115 |
let within = |t: Option<u64>, ms: u64| t.is_some_and(|t| now.saturating_sub(t) < ms); |
| 116 |
if a.paused { |
| 117 |
"paused" |
| 118 |
} else if a.blocked { |
| 119 |
"blocked" |
| 120 |
} else if within(a.status_ms, IDLE_AFTER_MS) { |
| 121 |
"busy" |
| 122 |
} else if a.waiting > 0 || within(a.last_request_ms, ONLINE_MS) { |
| 123 |
"online" |
| 124 |
} else if a.shims > 0 || within(a.last_request_ms, IDLE_AFTER_MS) { |
| 125 |
"idle" |
| 126 |
} else { |
| 127 |
"offline" |
| 128 |
} |
| 129 |
} |
| 130 |
|
| 131 |
|
| 132 |
pub async fn project(conn: &mut sqlx::SqliteConnection, e: &crate::event::Event) -> sqlx::Result<()> { |
| 133 |
let p = e.json(); |
| 134 |
let s = |k: &str| p.get(k).and_then(|v| v.as_str()).map(String::from); |
| 135 |
|
| 136 |
|
| 137 |
match e.kind.as_str() { |
| 138 |
"agent.presence" if s("agent").is_some() && s("state").is_some() => { |
| 139 |
sqlx::query("INSERT INTO presence (agent, state, since_ms) VALUES (?, ?, ?) ON CONFLICT(agent) DO UPDATE SET state = excluded.state, since_ms = excluded.since_ms") |
| 140 |
.bind(s("agent")) |
| 141 |
.bind(s("state")) |
| 142 |
.bind(e.at_ms as i64) |
| 143 |
.execute(&mut *conn) |
| 144 |
.await?; |
| 145 |
} |
| 146 |
"agent.status" if e.actor.starts_with("agent:") => { |
| 147 |
let agent = e.actor.strip_prefix("agent:").map(String::from); |
| 148 |
sqlx::query( |
| 149 |
"INSERT INTO presence (agent, state, since_ms, status_state, status_text, status_task, status_ms) VALUES (?, 'online', ?, ?, ?, ?, ?) \ |
| 150 |
ON CONFLICT(agent) DO UPDATE SET status_state = excluded.status_state, status_text = excluded.status_text, status_task = excluded.status_task, status_ms = excluded.status_ms", |
| 151 |
) |
| 152 |
.bind(agent) |
| 153 |
.bind(e.at_ms as i64) |
| 154 |
.bind(s("state")) |
| 155 |
.bind(s("text")) |
| 156 |
.bind(s("task")) |
| 157 |
.bind(e.at_ms as i64) |
| 158 |
.execute(&mut *conn) |
| 159 |
.await?; |
| 160 |
} |
| 161 |
"agent.session_started" if s("agent").is_some() && s("session").is_some() => { |
| 162 |
sqlx::query("INSERT OR IGNORE INTO sessions (id, agent, started_ms) VALUES (?, ?, ?)") |
| 163 |
.bind(s("session")) |
| 164 |
.bind(s("agent")) |
| 165 |
.bind(e.at_ms as i64) |
| 166 |
.execute(&mut *conn) |
| 167 |
.await?; |
| 168 |
} |
| 169 |
"agent.session_ended" => { |
| 170 |
sqlx::query("UPDATE sessions SET ended_ms = ? WHERE id = ?") |
| 171 |
.bind(e.at_ms as i64) |
| 172 |
.bind(s("session")) |
| 173 |
.execute(&mut *conn) |
| 174 |
.await?; |
| 175 |
} |
| 176 |
_ => {} |
| 177 |
} |
| 178 |
Ok(()) |
| 179 |
} |
| 180 |
|
| 181 |
#[cfg(test)] |
| 182 |
mod tests { |
| 183 |
use super::*; |
| 184 |
|
| 185 |
fn kinds(v: &[NewEvent]) -> Vec<String> { |
| 186 |
v.iter().map(|e| format!("{} {}", e.kind, e.payload.get("state").and_then(|s| s.as_str()).unwrap_or(""))).collect() |
| 187 |
} |
| 188 |
|
| 189 |
#[test] |
| 190 |
fn presence_moves_through_its_states_and_records_each_change() { |
| 191 |
let mut p = Presence::default(); |
| 192 |
let t = 1_000_000; |
| 193 |
let (s1, started) = p.request("claude", t, false); |
| 194 |
assert!(started.is_some()); |
| 195 |
assert_eq!(kinds(&p.tick(t)), ["agent.presence online"]); |
| 196 |
assert!(p.tick(t + 1).is_empty(), "no change, no event"); |
| 197 |
assert_eq!(p.state("claude", t + ONLINE_MS + 1), "idle"); |
| 198 |
assert_eq!(kinds(&p.tick(t + ONLINE_MS + 1)), ["agent.presence idle"]); |
| 199 |
p.status("claude", Some(t + ONLINE_MS + 2)); |
| 200 |
assert_eq!(p.state("claude", t + ONLINE_MS + 3), "busy"); |
| 201 |
p.blocked("claude", true); |
| 202 |
assert_eq!(p.state("claude", t + ONLINE_MS + 3), "blocked"); |
| 203 |
p.blocked("claude", false); |
| 204 |
p.status("claude", None); |
| 205 |
p.paused("claude", true); |
| 206 |
assert_eq!(p.state("claude", t + ONLINE_MS + 3), "paused"); |
| 207 |
p.paused("claude", false); |
| 208 |
let late = t + IDLE_AFTER_MS + 5; |
| 209 |
let ev = p.tick(late); |
| 210 |
assert_eq!(kinds(&ev), ["agent.presence offline", "agent.session_ended "]); |
| 211 |
let (s2, started) = p.request("claude", late + 1, false); |
| 212 |
assert!(started.is_some() && s2 != s1, "a new session after ten minutes' silence"); |
| 213 |
let (s3, none) = p.request("claude", late + 2, false); |
| 214 |
assert!(none.is_none() && s3 == s2); |
| 215 |
let (s4, again) = p.request("claude", late + 3, true); |
| 216 |
assert!(again.is_some() && s4 != s3, "brief(new_session: true)"); |
| 217 |
} |
| 218 |
|
| 219 |
#[test] |
| 220 |
fn a_waiting_or_connected_agent_is_not_offline() { |
| 221 |
let mut p = Presence::default(); |
| 222 |
p.request("kimi", 0, false); |
| 223 |
p.wait_started("kimi"); |
| 224 |
assert_eq!(p.state("kimi", IDLE_AFTER_MS * 3), "online"); |
| 225 |
p.wait_ended("kimi"); |
| 226 |
p.shim_connected("kimi"); |
| 227 |
assert_eq!(p.state("kimi", IDLE_AFTER_MS * 3), "idle"); |
| 228 |
p.shim_gone("kimi"); |
| 229 |
assert_eq!(p.state("kimi", IDLE_AFTER_MS * 3), "offline"); |
| 230 |
} |
| 231 |
} |