orch/store/
control.rs

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

orch/store/
control.rs
