orch/store/
control.rs

1use super::{Context, TaskStore};
2use chrono::{Duration, Utc};
3use serde::{Deserialize, Serialize};
4use sqlx::Row;
5use sqlx::sqlite::SqliteRow;
6
7/// Explicit column list for `SELECT` queries on the `control_messages` table.
8///
9/// Using `SELECT *` with `sqlx-sqlite`'s prepared statement cache can cause OOB
10/// panics when the schema changes (new migration columns) because the cached
11/// column metadata may not match the current table structure. Explicit columns
12/// prevent this mismatch.
13const CONTROL_MESSAGE_COLS: &str = "id, session_id, role, channel, channel_thread, \
14    content, summary, model, agent, input_tokens, output_tokens, tokens_used, cost_usd, created_at";
15
16/// A message in the control session conversation history.
17#[derive(Debug, Clone)]
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    #[expect(
61        clippy::too_many_arguments,
62        reason = "one argument per column of the message row"
63    )]
64    pub async fn insert_control_message(
65        &self,
66        session_id: &str,
67        role: &str,
68        channel: &str,
69        channel_thread: Option<&str>,
70        content: &str,
71        summary: Option<&str>,
72        model: Option<&str>,
73        agent: Option<&str>,
74        input_tokens: Option<i64>,
75        output_tokens: Option<i64>,
76        tokens_used: Option<i64>,
77        cost_usd: Option<f64>,
78    ) -> anyhow::Result<i64> {
79        let row = sqlx::query(
80        "INSERT INTO control_messages (session_id, role, channel, channel_thread, content, summary, model, agent, input_tokens, output_tokens, tokens_used, cost_usd)
81         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING id",
82    )
83    .bind(session_id)
84    .bind(role)
85    .bind(channel)
86    .bind(channel_thread)
87    .bind(content)
88    .bind(summary)
89    .bind(model)
90    .bind(agent)
91    .bind(input_tokens)
92    .bind(output_tokens)
93    .bind(tokens_used)
94    .bind(cost_usd)
95    .fetch_one(&self.pool)
96    .await?;
97        Ok(row.try_get("id")?)
98    }
99
100    /// List the most recent control messages for a session (chronological order).
101    ///
102    /// If `since` is provided, only messages created at or after that RFC3339/ISO8601
103    /// timestamp are returned. Use [`parse_since_duration`] to convert human-readable
104    /// durations like `"7d"` or `"24h"` into a timestamp string.
105    pub async fn list_control_messages(
106        &self,
107        session_id: &str,
108        since: Option<&str>,
109        limit: i64,
110    ) -> anyhow::Result<Vec<ControlMessage>> {
111        // Subquery: get the N most recent, then re-sort chronologically
112        let rows = if let Some(ts) = since {
113            sqlx::query(sqlx::AssertSqlSafe(format!(
114                "SELECT {CONTROL_MESSAGE_COLS} FROM (
115                SELECT {CONTROL_MESSAGE_COLS} FROM control_messages
116                WHERE session_id = ? AND created_at >= ?
117                ORDER BY created_at DESC, id DESC LIMIT ?
118            ) ORDER BY created_at ASC, id ASC"
119            )))
120            .bind(session_id)
121            .bind(ts)
122            .bind(limit)
123            .fetch_all(&self.pool)
124            .await?
125        } else {
126            sqlx::query(sqlx::AssertSqlSafe(format!(
127                "SELECT {CONTROL_MESSAGE_COLS} FROM (
128                SELECT {CONTROL_MESSAGE_COLS} FROM control_messages
129                WHERE session_id = ?
130                ORDER BY created_at DESC, id DESC LIMIT ?
131            ) ORDER BY created_at ASC, id ASC"
132            )))
133            .bind(session_id)
134            .bind(limit)
135            .fetch_all(&self.pool)
136            .await?
137        };
138        rows.iter()
139            .map(|r| {
140                Self::row_to_control_message(r).with_context(|| {
141                    format!(
142                        "parsing control_message row id={:?} session_id={:?}",
143                        r.try_get::<i64, _>("id").ok(),
144                        r.try_get::<String, _>("session_id").ok()
145                    )
146                })
147            })
148            .collect::<Result<Vec<_>, _>>()
149    }
150
151    /// Search control messages by content (LIKE match, most recent first).
152    ///
153    /// If `since` is provided, only messages created at or after that RFC3339/ISO8601
154    /// timestamp are returned. Use [`parse_since_duration`] to convert human-readable
155    /// durations like `"7d"` or `"24h"` into a timestamp string.
156    pub async fn search_control_messages(
157        &self,
158        session_id: &str,
159        query: &str,
160        since: Option<&str>,
161        limit: i64,
162    ) -> anyhow::Result<Vec<ControlMessage>> {
163        // Escape LIKE wildcards in user input
164        let escaped = query
165            .replace('\\', "\\\\")
166            .replace('%', "\\%")
167            .replace('_', "\\_");
168        let pattern = format!("%{escaped}%");
169        let rows = if let Some(ts) = since {
170            sqlx::query(sqlx::AssertSqlSafe(format!(
171                "SELECT {CONTROL_MESSAGE_COLS} FROM control_messages \
172                 WHERE session_id = ? AND content LIKE ? ESCAPE '\\' AND created_at >= ? \
173                 ORDER BY created_at DESC LIMIT ?"
174            )))
175            .bind(session_id)
176            .bind(&pattern)
177            .bind(ts)
178            .bind(limit)
179            .fetch_all(&self.pool)
180            .await?
181        } else {
182            sqlx::query(sqlx::AssertSqlSafe(format!(
183                "SELECT {CONTROL_MESSAGE_COLS} FROM control_messages \
184                 WHERE session_id = ? AND content LIKE ? ESCAPE '\\' \
185                 ORDER BY created_at DESC LIMIT ?"
186            )))
187            .bind(session_id)
188            .bind(&pattern)
189            .bind(limit)
190            .fetch_all(&self.pool)
191            .await?
192        };
193        rows.iter()
194            .map(|r| {
195                Self::row_to_control_message(r).with_context(|| {
196                    format!(
197                        "parsing control_message row id={:?} session_id={:?}",
198                        r.try_get::<i64, _>("id").ok(),
199                        r.try_get::<String, _>("session_id").ok()
200                    )
201                })
202            })
203            .collect::<Result<Vec<_>, _>>()
204    }
205
206    /// Get the N most recent assistant message summaries (chronological order).
207    pub async fn control_recent_summaries(
208        &self,
209        session_id: &str,
210        limit: i64,
211    ) -> anyhow::Result<Vec<String>> {
212        // Subquery: get the N most recent, then re-sort chronologically
213        let rows = sqlx::query(
214            "SELECT summary FROM (
215            SELECT summary, created_at, id FROM control_messages
216            WHERE session_id = ? AND summary IS NOT NULL
217            ORDER BY created_at DESC, id DESC LIMIT ?
218        ) ORDER BY created_at ASC, id ASC",
219        )
220        .bind(session_id)
221        .bind(limit)
222        .fetch_all(&self.pool)
223        .await?;
224        rows.iter()
225            .map(|r| r.try_get::<String, _>("summary"))
226            .collect::<Result<Vec<_>, _>>()
227            .map_err(Into::into)
228    }
229
230    fn row_to_control_message(row: &SqliteRow) -> Result<ControlMessage, sqlx::Error> {
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 {s:?}; expected a number followed by d, h, or m (e.g. 7d, 24h, 30m)"
385        );
386    };
387
388    let n: u64 = num_str.trim().parse().map_err(|_| {
389        anyhow::anyhow!(
390            "invalid number in --since {s:?}; expected a positive integer followed by d, h, or m"
391        )
392    })?;
393    if n == 0 {
394        anyhow::bail!("--since value must be greater than zero");
395    }
396
397    let secs: i64 = match unit {
398        'd' => n.checked_mul(86_400),
399        'h' => n.checked_mul(3_600),
400        'm' => n.checked_mul(60),
401        _ => anyhow::bail!("unrecognized duration unit '{unit}' (expected d, h, or m)"),
402    }
403    .ok_or_else(|| anyhow::anyhow!("--since value is too large"))?
404    .try_into()
405    .map_err(|_| anyhow::anyhow!("--since value is too large"))?;
406
407    Ok((Utc::now() - Duration::seconds(secs)).to_rfc3339_opts(chrono::SecondsFormat::Secs, true))
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    /// Helper: open an in-memory store and insert one control message via the public API.
415    async fn store_with_message() -> (TaskStore, i64) {
416        let store = TaskStore::open_memory().await.unwrap();
417        let id = store
418            .insert_control_message(
419                "sess1",
420                "user",
421                "cli",
422                None,
423                "hello world",
424                Some("a summary"),
425                Some("sonnet"),
426                Some("claude"),
427                Some(10),
428                Some(20),
429                Some(30),
430                Some(0.001),
431            )
432            .await
433            .unwrap();
434        (store, id)
435    }
436
437    /// All required fields must decode to their actual values, not silently default.
438    #[tokio::test]
439    async fn row_to_control_message_decodes_required_fields() {
440        let (store, _) = store_with_message().await;
441        let msgs = store
442            .list_control_messages("sess1", None, 10)
443            .await
444            .unwrap();
445        assert_eq!(msgs.len(), 1);
446        let m = &msgs[0];
447        assert_eq!(m.session_id, "sess1");
448        assert_eq!(m.role, "user");
449        assert_eq!(m.channel, "cli");
450        assert_eq!(m.content, "hello world");
451        assert_ne!(m.created_at, "");
452    }
453
454    /// Optional fields must be `Some(value)` when present and `None` when NULL.
455    #[tokio::test]
456    async fn row_to_control_message_decodes_optional_fields() {
457        let (store, _) = store_with_message().await;
458        let msgs = store
459            .list_control_messages("sess1", None, 10)
460            .await
461            .unwrap();
462        let m = &msgs[0];
463        assert_eq!(m.summary, Some("a summary".to_string()));
464        assert_eq!(m.model, Some("sonnet".to_string()));
465        assert_eq!(m.agent, Some("claude".to_string()));
466        assert_eq!(m.input_tokens, Some(10));
467        assert_eq!(m.output_tokens, Some(20));
468        assert_eq!(m.tokens_used, Some(30));
469        assert_eq!(m.cost_usd, Some(0.001));
470        assert_eq!(m.channel_thread, None);
471    }
472
473    /// Optional fields must be `None` when all nullable columns are NULL.
474    #[tokio::test]
475    async fn row_to_control_message_null_optional_fields_are_none() {
476        let store = TaskStore::open_memory().await.unwrap();
477        store
478            .insert_control_message(
479                "s",
480                "assistant",
481                "cli",
482                None,
483                "hi",
484                None,
485                None,
486                None,
487                None,
488                None,
489                None,
490                None,
491            )
492            .await
493            .unwrap();
494        let msgs = store.list_control_messages("s", None, 10).await.unwrap();
495        assert_eq!(msgs.len(), 1);
496        let m = &msgs[0];
497        assert!(m.summary.is_none());
498        assert!(m.model.is_none());
499        assert!(m.agent.is_none());
500        assert!(m.input_tokens.is_none());
501        assert!(m.output_tokens.is_none());
502        assert!(m.tokens_used.is_none());
503        assert!(m.cost_usd.is_none());
504        assert!(m.channel_thread.is_none());
505    }
506
507    /// Verify that decode errors from missing columns propagate via `?`.
508    ///
509    /// SQLite's loose type system coerces NULL to an empty string for TEXT columns,
510    /// so NULL-in-required-column does not produce a SQLx decode error. However,
511    /// a missing column (ColumnNotFound) does. This test confirms the `?` operator
512    /// actually propagates errors rather than swallowing them.
513    #[tokio::test]
514    async fn row_to_control_message_missing_column_propagates_error() {
515        let store = TaskStore::open_memory().await.unwrap();
516        // Select a row that omits the required `role` column entirely.
517        let rows = sqlx::query(
518            "SELECT 1 as id, 'sess' as session_id, 'cli' as channel, \
519             NULL as channel_thread, 'content' as content, NULL as summary, \
520             NULL as model, NULL as agent, NULL as input_tokens, NULL as output_tokens, \
521             NULL as tokens_used, NULL as cost_usd, '2024-01-01T00:00:00Z' as created_at",
522        )
523        .fetch_all(&store.pool)
524        .await
525        .unwrap();
526        assert_eq!(rows.len(), 1);
527        let result = TaskStore::row_to_control_message(&rows[0]);
528        assert!(
529            result.is_err(),
530            "expected Err when required column `role` is missing, got Ok"
531        );
532    }
533}
orch/store/control.rs

orch/store/
control.rs
