orch/store/
control.rs

1use super::*;
2use chrono::{Duration, Utc};
3use serde::{Deserialize, Serialize};
4use sqlx::Row;
5
6/// Explicit column list for `SELECT` queries on the `control_messages` table.
7///
8/// Using `SELECT *` with `sqlx-sqlite`'s prepared statement cache can cause OOB
9/// panics when the schema changes (new migration columns) because the cached
10/// column metadata may not match the current table structure. Explicit columns
11/// prevent this mismatch.
12const CONTROL_MESSAGE_COLS: &str = "id, session_id, role, channel, channel_thread, \
13    content, summary, model, agent, input_tokens, output_tokens, tokens_used, cost_usd, created_at";
14
15/// A message in the control session conversation history.
16#[derive(Debug, Clone)]
17#[allow(dead_code)] // fields used by channel integration (Phase 2)
18pub struct ControlMessage {
19    pub id: i64,
20    pub session_id: String,
21    pub role: String,
22    pub channel: String,
23    pub channel_thread: Option<String>,
24    pub content: String,
25    pub summary: Option<String>,
26    pub model: Option<String>,
27    pub agent: Option<String>,
28    pub input_tokens: Option<i64>,
29    pub output_tokens: Option<i64>,
30    pub tokens_used: Option<i64>,
31    pub cost_usd: Option<f64>,
32    pub created_at: String,
33}
34
35/// A single memory entry from a task attempt.
36#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
37pub struct MemoryEntry {
38    /// Attempt number (1-indexed)
39    pub attempt: u32,
40    /// Agent that made this attempt
41    pub agent: String,
42    /// Model used for this attempt
43    pub model: Option<String>,
44    /// Key learnings from this attempt
45    pub learnings: Vec<String>,
46    /// Error message if the attempt failed
47    pub error: Option<String>,
48    /// Files modified in this attempt
49    pub files_modified: Vec<String>,
50    /// Approach taken (from summary)
51    pub approach: String,
52    /// Timestamp of the attempt
53    pub timestamp: String,
54}
55
56impl TaskStore {
57    pub const DEFAULT_SESSION: &'static str = "default";
58
59    /// Insert a control session message.
60    #[allow(clippy::too_many_arguments)]
61    pub async fn insert_control_message(
62        &self,
63        session_id: &str,
64        role: &str,
65        channel: &str,
66        channel_thread: Option<&str>,
67        content: &str,
68        summary: Option<&str>,
69        model: Option<&str>,
70        agent: Option<&str>,
71        input_tokens: Option<i64>,
72        output_tokens: Option<i64>,
73        tokens_used: Option<i64>,
74        cost_usd: Option<f64>,
75    ) -> anyhow::Result<i64> {
76        let row = sqlx::query(
77        "INSERT INTO control_messages (session_id, role, channel, channel_thread, content, summary, model, agent, input_tokens, output_tokens, tokens_used, cost_usd)
78         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING id",
79    )
80    .bind(session_id)
81    .bind(role)
82    .bind(channel)
83    .bind(channel_thread)
84    .bind(content)
85    .bind(summary)
86    .bind(model)
87    .bind(agent)
88    .bind(input_tokens)
89    .bind(output_tokens)
90    .bind(tokens_used)
91    .bind(cost_usd)
92    .fetch_one(&self.pool)
93    .await?;
94        Ok(row.try_get("id")?)
95    }
96
97    /// List the most recent control messages for a session (chronological order).
98    ///
99    /// If `since` is provided, only messages created at or after that RFC3339/ISO8601
100    /// timestamp are returned. Use [`parse_since_duration`] to convert human-readable
101    /// durations like `"7d"` or `"24h"` into a timestamp string.
102    pub async fn list_control_messages(
103        &self,
104        session_id: &str,
105        since: Option<&str>,
106        limit: i64,
107    ) -> anyhow::Result<Vec<ControlMessage>> {
108        // Subquery: get the N most recent, then re-sort chronologically
109        let rows = if let Some(ts) = since {
110            sqlx::query(&format!(
111                "SELECT {CONTROL_MESSAGE_COLS} FROM (
112                SELECT {CONTROL_MESSAGE_COLS} FROM control_messages
113                WHERE session_id = ? AND created_at >= ?
114                ORDER BY created_at DESC, id DESC LIMIT ?
115            ) ORDER BY created_at ASC, id ASC"
116            ))
117            .bind(session_id)
118            .bind(ts)
119            .bind(limit)
120            .fetch_all(&self.pool)
121            .await?
122        } else {
123            sqlx::query(&format!(
124                "SELECT {CONTROL_MESSAGE_COLS} FROM (
125                SELECT {CONTROL_MESSAGE_COLS} FROM control_messages
126                WHERE session_id = ?
127                ORDER BY created_at DESC, id DESC LIMIT ?
128            ) ORDER BY created_at ASC, id ASC"
129            ))
130            .bind(session_id)
131            .bind(limit)
132            .fetch_all(&self.pool)
133            .await?
134        };
135        rows.iter()
136            .map(|r| {
137                Self::row_to_control_message(r).with_context(|| {
138                    format!(
139                        "parsing control_message row id={:?} session_id={:?}",
140                        r.try_get::<i64, _>("id").ok(),
141                        r.try_get::<String, _>("session_id").ok()
142                    )
143                })
144            })
145            .collect::<Result<Vec<_>, _>>()
146    }
147
148    /// Search control messages by content (LIKE match, most recent first).
149    ///
150    /// If `since` is provided, only messages created at or after that RFC3339/ISO8601
151    /// timestamp are returned. Use [`parse_since_duration`] to convert human-readable
152    /// durations like `"7d"` or `"24h"` into a timestamp string.
153    pub async fn search_control_messages(
154        &self,
155        session_id: &str,
156        query: &str,
157        since: Option<&str>,
158        limit: i64,
159    ) -> anyhow::Result<Vec<ControlMessage>> {
160        // Escape LIKE wildcards in user input
161        let escaped = query
162            .replace('\\', "\\\\")
163            .replace('%', "\\%")
164            .replace('_', "\\_");
165        let pattern = format!("%{escaped}%");
166        let rows = if let Some(ts) = since {
167            sqlx::query(&format!(
168                "SELECT {CONTROL_MESSAGE_COLS} FROM control_messages \
169                 WHERE session_id = ? AND content LIKE ? ESCAPE '\\' AND created_at >= ? \
170                 ORDER BY created_at DESC LIMIT ?"
171            ))
172            .bind(session_id)
173            .bind(&pattern)
174            .bind(ts)
175            .bind(limit)
176            .fetch_all(&self.pool)
177            .await?
178        } else {
179            sqlx::query(&format!(
180                "SELECT {CONTROL_MESSAGE_COLS} FROM control_messages \
181                 WHERE session_id = ? AND content LIKE ? ESCAPE '\\' \
182                 ORDER BY created_at DESC LIMIT ?"
183            ))
184            .bind(session_id)
185            .bind(&pattern)
186            .bind(limit)
187            .fetch_all(&self.pool)
188            .await?
189        };
190        rows.iter()
191            .map(|r| {
192                Self::row_to_control_message(r).with_context(|| {
193                    format!(
194                        "parsing control_message row id={:?} session_id={:?}",
195                        r.try_get::<i64, _>("id").ok(),
196                        r.try_get::<String, _>("session_id").ok()
197                    )
198                })
199            })
200            .collect::<Result<Vec<_>, _>>()
201    }
202
203    /// Get the N most recent assistant message summaries (chronological order).
204    pub async fn control_recent_summaries(
205        &self,
206        session_id: &str,
207        limit: i64,
208    ) -> anyhow::Result<Vec<String>> {
209        // Subquery: get the N most recent, then re-sort chronologically
210        let rows = sqlx::query(
211            "SELECT summary FROM (
212            SELECT summary, created_at, id FROM control_messages
213            WHERE session_id = ? AND summary IS NOT NULL
214            ORDER BY created_at DESC, id DESC LIMIT ?
215        ) ORDER BY created_at ASC, id ASC",
216        )
217        .bind(session_id)
218        .bind(limit)
219        .fetch_all(&self.pool)
220        .await?;
221        rows.iter()
222            .map(|r| r.try_get::<String, _>("summary"))
223            .collect::<Result<Vec<_>, _>>()
224            .map_err(Into::into)
225    }
226
227    fn row_to_control_message(
228        row: &sqlx::sqlite::SqliteRow,
229    ) -> Result<ControlMessage, sqlx::Error> {
230        use sqlx::Row;
231        Ok(ControlMessage {
232            id: row.try_get("id")?,
233            session_id: row.try_get("session_id")?,
234            role: row.try_get("role")?,
235            channel: row.try_get("channel")?,
236            channel_thread: row.try_get::<Option<String>, _>("channel_thread")?,
237            content: row.try_get("content")?,
238            summary: row.try_get::<Option<String>, _>("summary")?,
239            model: row.try_get::<Option<String>, _>("model")?,
240            agent: row.try_get::<Option<String>, _>("agent")?,
241            input_tokens: row.try_get::<Option<i64>, _>("input_tokens")?,
242            output_tokens: row.try_get::<Option<i64>, _>("output_tokens")?,
243            tokens_used: row.try_get::<Option<i64>, _>("tokens_used")?,
244            cost_usd: row.try_get::<Option<f64>, _>("cost_usd")?,
245            created_at: row.try_get("created_at")?,
246        })
247    }
248}
249
250/// Cost summary for a chat session.
251#[derive(Debug, Clone, Default)]
252pub struct ChatCostSummary {
253    /// Total number of messages (user + assistant)
254    pub total_messages: i64,
255    /// Total number of assistant messages (responses from LLM)
256    pub assistant_messages: i64,
257    /// Total tokens used across all messages
258    pub total_tokens: i64,
259    /// Input tokens from assistant messages.
260    pub total_input_tokens: i64,
261    /// Output tokens from assistant messages.
262    pub total_output_tokens: i64,
263    /// Total estimated cost in USD
264    pub total_cost_usd: f64,
265    /// Breakdown by model (model name -> (message_count, tokens, cost))
266    pub by_model: Vec<(String, i64, i64, f64)>,
267    /// The most frequently used model
268    pub primary_model: Option<String>,
269    /// The agent used (typically consistent per session)
270    pub primary_agent: Option<String>,
271}
272
273impl TaskStore {
274    /// Get cost summary for a chat session.
275    ///
276    /// Returns aggregated statistics including total messages, tokens, cost,
277    /// and breakdown by model.
278    pub async fn get_session_cost_summary(
279        &self,
280        session_id: &str,
281    ) -> anyhow::Result<ChatCostSummary> {
282        // Get total counts and cost
283        let row = sqlx::query(
284            "SELECT
285                COUNT(*) as total_messages,
286                COUNT(CASE WHEN role = 'assistant' THEN 1 END) as assistant_messages,
287                COALESCE(SUM(tokens_used), 0) as total_tokens,
288                COALESCE(SUM(input_tokens), 0) as total_input_tokens,
289                COALESCE(SUM(output_tokens), 0) as total_output_tokens,
290                COALESCE(SUM(cost_usd), 0.0) as total_cost_usd
291             FROM control_messages
292             WHERE session_id = ?",
293        )
294        .bind(session_id)
295        .fetch_one(&self.pool)
296        .await?;
297
298        let total_messages: i64 = row.try_get("total_messages")?;
299        let assistant_messages: i64 = row.try_get("assistant_messages")?;
300        let total_tokens: i64 = row.try_get("total_tokens")?;
301        let total_input_tokens: i64 = row.try_get("total_input_tokens")?;
302        let total_output_tokens: i64 = row.try_get("total_output_tokens")?;
303        let total_cost_usd: f64 = row.try_get("total_cost_usd")?;
304
305        // Get breakdown by model
306        let model_rows = sqlx::query(
307            "SELECT
308                COALESCE(model, 'unknown') as model_name,
309                COUNT(*) as message_count,
310                COALESCE(SUM(tokens_used), 0) as tokens,
311                COALESCE(SUM(cost_usd), 0.0) as cost
312             FROM control_messages
313             WHERE session_id = ? AND role = 'assistant'
314             GROUP BY model
315             ORDER BY message_count DESC",
316        )
317        .bind(session_id)
318        .fetch_all(&self.pool)
319        .await?;
320
321        let mut by_model = Vec::new();
322        for row in model_rows {
323            let model: String = row.try_get("model_name")?;
324            let count: i64 = row.try_get("message_count")?;
325            let tokens: i64 = row.try_get("tokens")?;
326            let cost: f64 = row.try_get("cost")?;
327            by_model.push((model, count, tokens, cost));
328        }
329
330        // Determine primary model (most used)
331        let primary_model = by_model.first().map(|(m, _, _, _)| m.clone());
332
333        // Get primary agent (most used)
334        let agent_row = sqlx::query(
335            "SELECT COALESCE(agent, 'unknown') as agent_name
336             FROM control_messages
337             WHERE session_id = ? AND role = 'assistant' AND agent IS NOT NULL
338             GROUP BY agent
339             ORDER BY COUNT(*) DESC
340             LIMIT 1",
341        )
342        .bind(session_id)
343        .fetch_optional(&self.pool)
344        .await?;
345
346        let primary_agent = agent_row.and_then(|r| r.try_get::<String, _>("agent_name").ok());
347
348        Ok(ChatCostSummary {
349            total_messages,
350            assistant_messages,
351            total_tokens,
352            total_input_tokens,
353            total_output_tokens,
354            total_cost_usd,
355            by_model,
356            primary_model,
357            primary_agent,
358        })
359    }
360}
361
362/// Parse a human-readable duration string into an RFC3339 timestamp representing
363/// `now - duration`.
364///
365/// Supported suffixes:
366/// - `d` — days (e.g. `"7d"`)
367/// - `h` — hours (e.g. `"24h"`)
368/// - `m` — minutes (e.g. `"30m"`)
369///
370/// Returns an ISO8601/RFC3339 string suitable for comparison against
371/// `control_messages.created_at` (stored as UTC RFC3339).
372///
373/// Returns `Err` if the format is unrecognised or the number overflows.
374pub fn parse_since_duration(s: &str) -> anyhow::Result<String> {
375    let s = s.trim();
376    let (num_str, unit) = if let Some(n) = s.strip_suffix('d') {
377        (n, 'd')
378    } else if let Some(n) = s.strip_suffix('h') {
379        (n, 'h')
380    } else if let Some(n) = s.strip_suffix('m') {
381        (n, 'm')
382    } else {
383        anyhow::bail!(
384            "unrecognised --since format {:?}; expected a number followed by d, h, or m (e.g. 7d, 24h, 30m)",
385            s
386        );
387    };
388
389    let n: u64 = num_str.trim().parse().map_err(|_| {
390        anyhow::anyhow!(
391            "invalid number in --since {:?}; expected a positive integer followed by d, h, or m",
392            s
393        )
394    })?;
395    if n == 0 {
396        anyhow::bail!("--since value must be greater than zero");
397    }
398
399    let secs: i64 = match unit {
400        'd' => n.checked_mul(86_400),
401        'h' => n.checked_mul(3_600),
402        'm' => n.checked_mul(60),
403        _ => anyhow::bail!(
404            "unrecognized duration unit '{}' (expected d, h, or m)",
405            unit
406        ),
407    }
408    .ok_or_else(|| anyhow::anyhow!("--since value is too large"))?
409    .try_into()
410    .map_err(|_| anyhow::anyhow!("--since value is too large"))?;
411
412    Ok((Utc::now() - Duration::seconds(secs)).to_rfc3339_opts(chrono::SecondsFormat::Secs, true))
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    /// Helper: open an in-memory store and insert one control message via the public API.
420    async fn store_with_message() -> (TaskStore, i64) {
421        let store = TaskStore::open_memory().await.unwrap();
422        let id = store
423            .insert_control_message(
424                "sess1",
425                "user",
426                "cli",
427                None,
428                "hello world",
429                Some("a summary"),
430                Some("sonnet"),
431                Some("claude"),
432                Some(10),
433                Some(20),
434                Some(30),
435                Some(0.001),
436            )
437            .await
438            .unwrap();
439        (store, id)
440    }
441
442    /// All required fields must decode to their actual values, not silently default.
443    #[tokio::test]
444    async fn row_to_control_message_decodes_required_fields() {
445        let (store, _) = store_with_message().await;
446        let msgs = store
447            .list_control_messages("sess1", None, 10)
448            .await
449            .unwrap();
450        assert_eq!(msgs.len(), 1);
451        let m = &msgs[0];
452        assert_eq!(m.session_id, "sess1");
453        assert_eq!(m.role, "user");
454        assert_eq!(m.channel, "cli");
455        assert_eq!(m.content, "hello world");
456        assert!(!m.created_at.is_empty());
457    }
458
459    /// Optional fields must be `Some(value)` when present and `None` when NULL.
460    #[tokio::test]
461    async fn row_to_control_message_decodes_optional_fields() {
462        let (store, _) = store_with_message().await;
463        let msgs = store
464            .list_control_messages("sess1", None, 10)
465            .await
466            .unwrap();
467        let m = &msgs[0];
468        assert_eq!(m.summary, Some("a summary".to_string()));
469        assert_eq!(m.model, Some("sonnet".to_string()));
470        assert_eq!(m.agent, Some("claude".to_string()));
471        assert_eq!(m.input_tokens, Some(10));
472        assert_eq!(m.output_tokens, Some(20));
473        assert_eq!(m.tokens_used, Some(30));
474        assert_eq!(m.cost_usd, Some(0.001));
475        assert_eq!(m.channel_thread, None);
476    }
477
478    /// Optional fields must be `None` when all nullable columns are NULL.
479    #[tokio::test]
480    async fn row_to_control_message_null_optional_fields_are_none() {
481        let store = TaskStore::open_memory().await.unwrap();
482        store
483            .insert_control_message(
484                "s",
485                "assistant",
486                "cli",
487                None,
488                "hi",
489                None,
490                None,
491                None,
492                None,
493                None,
494                None,
495                None,
496            )
497            .await
498            .unwrap();
499        let msgs = store.list_control_messages("s", None, 10).await.unwrap();
500        assert_eq!(msgs.len(), 1);
501        let m = &msgs[0];
502        assert!(m.summary.is_none());
503        assert!(m.model.is_none());
504        assert!(m.agent.is_none());
505        assert!(m.input_tokens.is_none());
506        assert!(m.output_tokens.is_none());
507        assert!(m.tokens_used.is_none());
508        assert!(m.cost_usd.is_none());
509        assert!(m.channel_thread.is_none());
510    }
511
512    /// Verify that decode errors from missing columns propagate via `?`.
513    ///
514    /// SQLite's loose type system coerces NULL to an empty string for TEXT columns,
515    /// so NULL-in-required-column does not produce a SQLx decode error. However,
516    /// a missing column (ColumnNotFound) does. This test confirms the `?` operator
517    /// actually propagates errors rather than swallowing them.
518    #[tokio::test]
519    async fn row_to_control_message_missing_column_propagates_error() {
520        let store = TaskStore::open_memory().await.unwrap();
521        // Select a row that omits the required `role` column entirely.
522        let rows = sqlx::query(
523            "SELECT 1 as id, 'sess' as session_id, 'cli' as channel, \
524             NULL as channel_thread, 'content' as content, NULL as summary, \
525             NULL as model, NULL as agent, NULL as input_tokens, NULL as output_tokens, \
526             NULL as tokens_used, NULL as cost_usd, '2024-01-01T00:00:00Z' as created_at",
527        )
528        .fetch_all(&store.pool)
529        .await
530        .unwrap();
531        assert_eq!(rows.len(), 1);
532        let result = TaskStore::row_to_control_message(&rows[0]);
533        assert!(
534            result.is_err(),
535            "expected Err when required column `role` is missing, got Ok"
536        );
537    }
538}
orch/store/control.rs

orch/store/
control.rs
