dk-protocol 0.2.77

dkod gRPC protocol definitions and server/client implementations
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use chrono::{DateTime, TimeDelta, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::time::Duration;
use uuid::Uuid;

pub type SessionId = Uuid;

pub struct AgentSession {
    pub id: SessionId,
    pub agent_id: String,
    pub codebase: String,
    pub intent: String,
    pub codebase_version: String,
    pub created_at: DateTime<Utc>,
    pub last_active: DateTime<Utc>,
}

/// Snapshot of a session's identity info, saved when a session expires or
/// is explicitly removed, allowing a new CONNECT to resume it.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SessionSnapshot {
    pub agent_id: String,
    pub codebase: String,
    pub intent: String,
    pub codebase_version: String,
}

pub struct SessionManager {
    sessions: DashMap<SessionId, AgentSession>,
    timeout: Duration,
    /// Snapshots of expired session workspaces for resume support.
    snapshots: DashMap<SessionId, SessionSnapshot>,
}

impl SessionManager {
    pub fn new(timeout: Duration) -> Self {
        Self {
            sessions: DashMap::new(),
            timeout,
            snapshots: DashMap::new(),
        }
    }

    pub fn create_session(
        &self,
        agent_id: String,
        codebase: String,
        intent: String,
        codebase_version: String,
    ) -> SessionId {
        let id = Uuid::new_v4();
        let now = Utc::now();
        self.sessions.insert(
            id,
            AgentSession {
                id,
                agent_id,
                codebase,
                intent,
                codebase_version,
                created_at: now,
                last_active: now,
            },
        );
        id
    }

    pub fn get_session(&self, id: &SessionId) -> Option<AgentSession> {
        let entry = self.sessions.get(id)?;
        let elapsed = Utc::now().signed_duration_since(entry.last_active);
        let timeout = TimeDelta::from_std(self.timeout).unwrap_or(TimeDelta::MAX);
        if elapsed > timeout {
            drop(entry);
            self.sessions.remove(id);
            return None;
        }
        Some(AgentSession {
            id: entry.id,
            agent_id: entry.agent_id.clone(),
            codebase: entry.codebase.clone(),
            intent: entry.intent.clone(),
            codebase_version: entry.codebase_version.clone(),
            created_at: entry.created_at,
            last_active: entry.last_active,
        })
    }

    pub fn touch_session(&self, id: &SessionId) -> bool {
        if let Some(mut entry) = self.sessions.get_mut(id) {
            entry.last_active = Utc::now();
            true
        } else {
            false
        }
    }

    pub fn remove_session(&self, id: &SessionId) -> bool {
        self.sessions.remove(id).is_some()
    }

    /// Save a snapshot of a session for later resume.
    pub fn save_snapshot(&self, id: &SessionId, snapshot: SessionSnapshot) {
        self.snapshots.insert(*id, snapshot);
    }

    /// Retrieve and remove a saved session snapshot.
    pub fn take_snapshot(&self, id: &SessionId) -> Option<SessionSnapshot> {
        self.snapshots.remove(id).map(|(_, snap)| snap)
    }

    pub fn cleanup_expired(&self) {
        let now = Utc::now();
        let timeout = TimeDelta::from_std(self.timeout).unwrap_or(TimeDelta::MAX);
        let mut expired = Vec::new();
        self.sessions.retain(|id, session| {
            let alive = now.signed_duration_since(session.last_active) <= timeout;
            if !alive {
                expired.push((*id, SessionSnapshot {
                    agent_id: session.agent_id.clone(),
                    codebase: session.codebase.clone(),
                    intent: session.intent.clone(),
                    codebase_version: session.codebase_version.clone(),
                }));
            }
            alive
        });
        for (id, snap) in expired {
            self.snapshots.insert(id, snap);
        }
    }
}