agent_trace/runtime/
session.rs1use crate::types::Actor;
2use anyhow::Result;
3use chrono::{DateTime, Duration, Utc};
4use serde::{Deserialize, Serialize};
5use std::path::Path;
6
7pub const LOCK_FILE: &str = ".agent-trace/locks/agent-lock.toml";
8const STALE_TIMEOUT_MINUTES: i64 = 30;
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
11pub struct AgentSession {
12 pub name: String,
13 pub session_id: String,
14 pub transport: String,
15 pub started_at: String,
16 pub last_heartbeat: String,
17}
18
19pub fn new_session_id() -> String {
20 let now = Utc::now();
21 format!(
22 "{}-{:03}",
23 now.format("%Y%m%d-%H%M%S"),
24 now.timestamp_subsec_millis()
25 )
26}
27
28impl AgentSession {
29 fn now_rfc3339() -> String {
30 Utc::now().to_rfc3339()
31 }
32
33 fn parse_ts(ts: &str) -> Option<DateTime<Utc>> {
34 chrono::DateTime::parse_from_rfc3339(ts)
35 .ok()
36 .map(|dt| dt.with_timezone(&Utc))
37 }
38
39 pub fn is_stale(&self) -> bool {
40 let Some(last) = Self::parse_ts(&self.last_heartbeat) else {
41 return true;
42 };
43 Utc::now().signed_duration_since(last) > Duration::minutes(STALE_TIMEOUT_MINUTES)
44 }
45
46 pub fn refresh_heartbeat(&mut self) {
47 self.last_heartbeat = Self::now_rfc3339();
48 }
49}
50
51#[derive(Debug, Serialize, Deserialize)]
52struct LockFile {
53 agent: AgentSession,
54}
55
56fn lock_path(store_root: &Path) -> std::path::PathBuf {
57 store_root.join(LOCK_FILE)
58}
59
60fn write_lock(store_root: &Path, session: &AgentSession) -> Result<()> {
61 let path = lock_path(store_root);
62 std::fs::create_dir_all(path.parent().expect("lock parent"))?;
63 let payload = toml::to_string(&LockFile {
64 agent: session.clone(),
65 })?;
66 std::fs::write(path, payload)?;
67 Ok(())
68}
69
70pub fn load_session(store_root: &Path) -> Option<AgentSession> {
71 let path = lock_path(store_root);
72 let content = std::fs::read_to_string(path).ok()?;
73 let parsed: LockFile = toml::from_str(&content).ok()?;
74 Some(parsed.agent)
75}
76
77pub fn start_session(store_root: &Path, name: &str, transport: &str) -> Result<AgentSession> {
78 if let Some(existing) = load_session(store_root) {
79 if existing.is_stale() {
80 crate::session_recap::maybe_recap_prior_session(store_root, &existing)?;
81 remove_session(store_root)?;
82 }
83 }
84
85 let session = AgentSession {
86 name: name.to_string(),
87 session_id: new_session_id(),
88 transport: transport.to_string(),
89 started_at: AgentSession::now_rfc3339(),
90 last_heartbeat: AgentSession::now_rfc3339(),
91 };
92 write_lock(store_root, &session)?;
93 Ok(session)
94}
95
96pub fn touch_session(store_root: &Path, expected_name: &str) -> Result<()> {
97 if let Some(mut session) = load_session(store_root) {
98 if session.name == expected_name && !session.is_stale() {
99 session.refresh_heartbeat();
100 write_lock(store_root, &session)?;
101 }
102 }
103 Ok(())
104}
105
106pub fn remove_session(store_root: &Path) -> Result<()> {
107 let path = lock_path(store_root);
108 if path.exists() {
109 std::fs::remove_file(path)?;
110 }
111 Ok(())
112}
113
114pub fn session_id_for_store(store_root: &Path) -> Option<String> {
116 load_session(store_root)
117 .filter(|s| !s.is_stale())
118 .map(|s| s.session_id)
119}
120
121pub fn session_id_for_actor(store_root: &Path, actor: &Actor) -> Option<String> {
122 let Actor::Agent { name } = actor else {
123 return None;
124 };
125 let session = load_session(store_root)?;
126 if session.name != *name || session.is_stale() {
127 return None;
128 }
129 Some(session.session_id)
130}
131
132pub struct AgentState {
139 pub cli_agent: Option<String>,
140}
141
142impl AgentState {
143 pub fn new(cli_agent: Option<String>) -> Self {
144 Self { cli_agent }
145 }
146
147 pub fn current_actor(&self, store_root: &Path) -> Actor {
148 if let Some(session) = load_session(store_root) {
149 if session.is_stale() {
150 let _ = crate::session_recap::maybe_recap_prior_session(store_root, &session);
151 let _ = remove_session(store_root);
152 } else {
153 return Actor::Agent { name: session.name };
154 }
155 }
156
157 if let Some(name) = &self.cli_agent {
158 return Actor::Agent { name: name.clone() };
159 }
160
161 Actor::User
162 }
163}
164
165#[cfg(test)]
166mod tests {
167 use super::*;
168 use tempfile::TempDir;
169
170 #[test]
171 fn stale_session_is_ignored() {
172 let tmp = TempDir::new().unwrap();
173 let root = tmp.path();
174 std::fs::create_dir_all(root.join(".agent-trace/locks")).unwrap();
175 std::fs::write(
176 root.join(LOCK_FILE),
177 "[agent]\nname=\"bot\"\nsession_id=\"s1\"\ntransport=\"cli\"\nstarted_at=\"2020-01-01T00:00:00Z\"\nlast_heartbeat=\"2020-01-01T00:00:00Z\"\n",
178 )
179 .unwrap();
180
181 let state = AgentState::new(None);
182 assert_eq!(state.current_actor(root), Actor::User);
183 }
184
185 #[test]
186 fn start_session_writes_metadata() {
187 let tmp = TempDir::new().unwrap();
188 let root = tmp.path();
189 let s = start_session(root, "worker", "mcp").unwrap();
190 assert_eq!(s.name, "worker");
191 assert_eq!(s.transport, "mcp");
192 let loaded = load_session(root).unwrap();
193 assert_eq!(loaded.name, "worker");
194 assert!(!loaded.session_id.is_empty());
195 }
196}