dk_protocol/
session.rs

1use chrono::{DateTime, TimeDelta, Utc};
2use dashmap::DashMap;
3use serde::{Deserialize, Serialize};
4use std::time::Duration;
5use uuid::Uuid;
6
7pub type SessionId = Uuid;
8
9pub struct AgentSession {
10    pub id: SessionId,
11    pub agent_id: String,
12    pub codebase: String,
13    pub intent: String,
14    pub codebase_version: String,
15    pub created_at: DateTime<Utc>,
16    pub last_active: DateTime<Utc>,
17}
18
19/// Snapshot of a session's identity info, saved when a session expires or
20/// is explicitly removed, allowing a new CONNECT to resume it.
21#[derive(Clone, Debug, Serialize, Deserialize)]
22pub struct SessionSnapshot {
23    pub agent_id: String,
24    pub codebase: String,
25    pub intent: String,
26    pub codebase_version: String,
27}
28
29pub struct SessionManager {
30    sessions: DashMap<SessionId, AgentSession>,
31    timeout: Duration,
32    /// Snapshots of expired session workspaces for resume support.
33    snapshots: DashMap<SessionId, SessionSnapshot>,
34}
35
36impl SessionManager {
37    pub fn new(timeout: Duration) -> Self {
38        Self {
39            sessions: DashMap::new(),
40            timeout,
41            snapshots: DashMap::new(),
42        }
43    }
44
45    pub fn create_session(
46        &self,
47        agent_id: String,
48        codebase: String,
49        intent: String,
50        codebase_version: String,
51    ) -> SessionId {
52        let id = Uuid::new_v4();
53        let now = Utc::now();
54        self.sessions.insert(
55            id,
56            AgentSession {
57                id,
58                agent_id,
59                codebase,
60                intent,
61                codebase_version,
62                created_at: now,
63                last_active: now,
64            },
65        );
66        id
67    }
68
69    pub fn get_session(&self, id: &SessionId) -> Option<AgentSession> {
70        let entry = self.sessions.get(id)?;
71        let elapsed = Utc::now().signed_duration_since(entry.last_active);
72        let timeout = TimeDelta::from_std(self.timeout).unwrap_or(TimeDelta::MAX);
73        if elapsed > timeout {
74            drop(entry);
75            self.sessions.remove(id);
76            return None;
77        }
78        Some(AgentSession {
79            id: entry.id,
80            agent_id: entry.agent_id.clone(),
81            codebase: entry.codebase.clone(),
82            intent: entry.intent.clone(),
83            codebase_version: entry.codebase_version.clone(),
84            created_at: entry.created_at,
85            last_active: entry.last_active,
86        })
87    }
88
89    pub fn touch_session(&self, id: &SessionId) -> bool {
90        if let Some(mut entry) = self.sessions.get_mut(id) {
91            entry.last_active = Utc::now();
92            true
93        } else {
94            false
95        }
96    }
97
98    pub fn remove_session(&self, id: &SessionId) -> bool {
99        self.sessions.remove(id).is_some()
100    }
101
102    /// Save a snapshot of a session for later resume.
103    pub fn save_snapshot(&self, id: &SessionId, snapshot: SessionSnapshot) {
104        self.snapshots.insert(*id, snapshot);
105    }
106
107    /// Retrieve and remove a saved session snapshot.
108    pub fn take_snapshot(&self, id: &SessionId) -> Option<SessionSnapshot> {
109        self.snapshots.remove(id).map(|(_, snap)| snap)
110    }
111
112    pub fn cleanup_expired(&self) {
113        let now = Utc::now();
114        let timeout = TimeDelta::from_std(self.timeout).unwrap_or(TimeDelta::MAX);
115        let mut expired = Vec::new();
116        self.sessions.retain(|id, session| {
117            let alive = now.signed_duration_since(session.last_active) <= timeout;
118            if !alive {
119                expired.push((*id, SessionSnapshot {
120                    agent_id: session.agent_id.clone(),
121                    codebase: session.codebase.clone(),
122                    intent: session.intent.clone(),
123                    codebase_version: session.codebase_version.clone(),
124                }));
125            }
126            alive
127        });
128        for (id, snap) in expired {
129            self.snapshots.insert(id, snap);
130        }
131    }
132}
dk_protocol/session.rs

dk_protocol/
session.rs
