Skip to main content

agent_trace/runtime/
session.rs

1use 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
114/// Returns the active session ID from the lock file, if non-stale.
115pub 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
132/// Resolves the current actor for a command/session.
133///
134/// Priority:
135/// 1) active lock file from `agent-trace connect`
136/// 2) explicit CLI flag (`--agent` / `--actor`)
137/// 3) fallback to user
138pub 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}