teamcraft/services/
assistant_threads.rs

1//! The assistant's conversation: one thread per user that survives page
2//! navigation and reloads.
3//!
4//! What is stored: the text of each question and of each answer, the quick
5//! action, and *where* the question was asked (the label of the resolved page
6//! context, for example `PRT-000009/A · bom`). **Not** stored: tool results,
7//! the model's reasoning, the page context itself. A follow-up replays the
8//! stored text to the model and calls the tools again with the user's
9//! *current* rights, so nothing they lost access to comes back through a
10//! tool. Only the owner reads a thread: every query here is keyed by the
11//! user id of the actor, there is no function that takes a thread id from a
12//! client. Threads expire after `TEAMCRAFT_ASSISTANT_RETENTION_DAYS` of
13//! inactivity and go with the user.
14
15use chrono::{DateTime, Utc};
16use serde::Serialize;
17use sqlx::{FromRow, PgConnection, PgPool};
18use uuid::Uuid;
19
20use crate::{auth::Actor, error::AppResult, services::audit};
21
22/// Default retention of an idle thread (days).
23pub const DEFAULT_RETENTION_DAYS: u32 = 30;
24/// Most messages replayed to the model on a follow-up (the newest ones).
25pub const MAX_HISTORY_MESSAGES: usize = 16;
26/// Most characters replayed per message.
27pub const MAX_HISTORY_CHARS: usize = 8_000;
28/// Most messages kept in a thread; the oldest go first.
29pub const MAX_STORED_MESSAGES: i64 = 200;
30
31/// One stored message, as the thread view shows it.
32#[derive(Debug, Clone, FromRow, Serialize)]
33pub struct StoredMessage {
34    /// `user` or `assistant`.
35    pub role: String,
36    pub text: String,
37    /// The quick action of a user message (`ask` for a free question).
38    pub action: String,
39    /// Kind and label of the page the user asked on (user messages).
40    pub context_kind: String,
41    pub context_label: String,
42    /// The run that produced this answer had read restricted content (see
43    /// `assistant::McpTools`); a follow-up starts tainted.
44    pub tainted: bool,
45    pub created_at: DateTime<Utc>,
46}
47
48/// The thread of a user.
49#[derive(Debug, Clone, Serialize)]
50pub struct Thread {
51    pub messages: Vec<StoredMessage>,
52    pub retention_days: u32,
53}
54
55/// Deletes the idle threads of everybody (the periodic clean-up) and
56/// returns how many went.
57pub async fn purge_expired(pool: &PgPool, retention_days: u32) -> AppResult<u64> {
58    let done = sqlx::query(
59        "DELETE FROM assistant_threads WHERE updated_at < now() - make_interval(days => $1)",
60    )
61    .bind(retention_days as i32)
62    .execute(pool)
63    .await?;
64    Ok(done.rows_affected())
65}
66
67/// The user's thread, newest message last. An expired thread is deleted and
68/// reads as empty.
69pub async fn load(pool: &PgPool, me: &Actor, retention_days: u32) -> AppResult<Thread> {
70    sqlx::query(
71        "DELETE FROM assistant_threads
72         WHERE user_id = $1 AND updated_at < now() - make_interval(days => $2)",
73    )
74    .bind(me.id)
75    .bind(retention_days as i32)
76    .execute(pool)
77    .await?;
78    let messages = sqlx::query_as::<_, StoredMessage>(
79        "SELECT m.role, m.text, m.action, m.context_kind, m.context_label, m.tainted, m.created_at
80         FROM assistant_messages m
81         JOIN assistant_threads t ON t.id = m.thread_id
82         WHERE t.user_id = $1
83         ORDER BY m.seq",
84    )
85    .bind(me.id)
86    .fetch_all(pool)
87    .await?;
88    Ok(Thread {
89        messages,
90        retention_days,
91    })
92}
93
94/// What a follow-up replays to the model: the newest messages, each cut.
95pub fn history_of(thread: &Thread) -> Vec<(String, String)> {
96    let skip = thread.messages.len().saturating_sub(MAX_HISTORY_MESSAGES);
97    let mut out: Vec<(String, String)> = thread
98        .messages
99        .iter()
100        .skip(skip)
101        .map(|m| {
102            let text: String = m.text.chars().take(MAX_HISTORY_CHARS).collect();
103            (m.role.clone(), text)
104        })
105        .collect();
106    // the conversation must start with the user
107    while out.first().is_some_and(|(role, _)| role != "user") {
108        out.remove(0);
109    }
110    out
111}
112
113/// Whether an earlier answer of the thread came from a tainted run.
114pub fn was_tainted(thread: &Thread) -> bool {
115    thread.messages.iter().any(|m| m.tainted)
116}
117
118/// A question and its answer to append.
119pub struct Exchange<'a> {
120    pub action: &'a str,
121    /// What the user typed, or the label of the quick action.
122    pub question: &'a str,
123    pub context_kind: &'a str,
124    pub context_label: &'a str,
125    pub answer: &'a str,
126    pub tainted: bool,
127}
128
129/// Appends a question and its answer in one transaction (the thread is
130/// created on the first exchange). Questions that were cancelled or failed
131/// are never stored.
132pub async fn append(conn: &mut PgConnection, me: &Actor, x: &Exchange<'_>) -> AppResult<()> {
133    let thread: Uuid = sqlx::query_scalar(
134        "INSERT INTO assistant_threads (id, user_id) VALUES ($1, $2)
135         ON CONFLICT (user_id) DO UPDATE SET updated_at = now()
136         RETURNING id",
137    )
138    .bind(Uuid::new_v4())
139    .bind(me.id)
140    .fetch_one(&mut *conn)
141    .await?;
142    // serialise concurrent questions of the same user (two tabs)
143    sqlx::query("SELECT id FROM assistant_threads WHERE id = $1 FOR UPDATE")
144        .bind(thread)
145        .fetch_one(&mut *conn)
146        .await?;
147    let next: i32 = sqlx::query_scalar(
148        "SELECT COALESCE(MAX(seq), 0) + 1 FROM assistant_messages WHERE thread_id = $1",
149    )
150    .bind(thread)
151    .fetch_one(&mut *conn)
152    .await?;
153    let rows = [
154        ("user", x.question, x.context_kind, x.context_label, false),
155        ("assistant", x.answer, "", "", x.tainted),
156    ];
157    for (n, (role, text, kind, label, tainted)) in rows.into_iter().enumerate() {
158        sqlx::query(
159            "INSERT INTO assistant_messages
160                 (thread_id, seq, role, text, action, context_kind, context_label, tainted)
161             VALUES ($1, $2, $3, $4, $5, $6, $7, $8)",
162        )
163        .bind(thread)
164        .bind(next + n as i32)
165        .bind(role)
166        .bind(text)
167        .bind(if role == "user" { x.action } else { "" })
168        .bind(kind)
169        .bind(label)
170        .bind(tainted)
171        .execute(&mut *conn)
172        .await?;
173    }
174    sqlx::query(
175        "DELETE FROM assistant_messages
176         WHERE thread_id = $1
177           AND seq <= (SELECT MAX(seq) FROM assistant_messages WHERE thread_id = $1) - $2",
178    )
179    .bind(thread)
180    .bind(MAX_STORED_MESSAGES)
181    .execute(&mut *conn)
182    .await?;
183    Ok(())
184}
185
186/// "New conversation": deletes the user's thread.
187///
188/// # Invariants
189///
190/// * **Permission:** the actor's own thread only (there is one per user).
191/// * **Locks:** none beyond the row locks of the `DELETE`.
192/// * **Audit:** `assistant.thread_reset`, only when a thread existed; the
193///   record has no content of the conversation (`tests/assistant_threads.rs`).
194/// * **Notifications / webhooks:** none.
195/// * **Atomic:** one transaction; without a thread nothing happens and no
196///   error is returned.
197pub async fn reset(pool: &PgPool, me: &Actor) -> AppResult<()> {
198    let mut tx = pool.begin().await?;
199    let done = sqlx::query("DELETE FROM assistant_threads WHERE user_id = $1")
200        .bind(me.id)
201        .execute(&mut *tx)
202        .await?;
203    if done.rows_affected() > 0 {
204        audit::record(
205            &mut tx,
206            Some(me),
207            "assistant.thread_reset",
208            "assistant",
209            None,
210            "Assistant: conversation deleted",
211            serde_json::json!({}),
212        )
213        .await?;
214    }
215    tx.commit().await?;
216    Ok(())
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    fn msg(role: &str, text: &str, tainted: bool) -> StoredMessage {
224        StoredMessage {
225            role: role.into(),
226            text: text.into(),
227            action: "ask".into(),
228            context_kind: String::new(),
229            context_label: String::new(),
230            tainted,
231            created_at: Utc::now(),
232        }
233    }
234
235    #[test]
236    fn history_is_capped_and_starts_with_the_user() {
237        let mut messages = Vec::new();
238        for n in 0..20 {
239            messages.push(msg("user", &format!("q{n}"), false));
240            messages.push(msg("assistant", &"a".repeat(MAX_HISTORY_CHARS + 50), false));
241        }
242        let thread = Thread {
243            messages,
244            retention_days: 30,
245        };
246        let h = history_of(&thread);
247        assert_eq!(h.len(), MAX_HISTORY_MESSAGES);
248        assert_eq!(h[0].0, "user");
249        assert_eq!(h[0].1, "q12");
250        assert!(h
251            .iter()
252            .all(|(_, t)| t.chars().count() <= MAX_HISTORY_CHARS));
253        // an odd cut drops the leading answer
254        let thread = Thread {
255            messages: vec![msg("assistant", "x", false), msg("user", "y", false)],
256            retention_days: 30,
257        };
258        assert_eq!(
259            history_of(&thread),
260            vec![("user".to_string(), "y".to_string())]
261        );
262    }
263
264    #[test]
265    fn taint_is_remembered() {
266        let mut thread = Thread {
267            messages: vec![msg("user", "q", false), msg("assistant", "a", false)],
268            retention_days: 30,
269        };
270        assert!(!was_tainted(&thread));
271        thread.messages[1].tainted = true;
272        assert!(was_tainted(&thread));
273    }
274}
teamcraft/services/assistant_threads.rs

teamcraft/services/
assistant_threads.rs
