1use anyhow::Result;
2use rusqlite::{params, Connection, OptionalExtension};
3
4#[derive(Debug, Clone, PartialEq)]
5pub struct LatestSessionMemorySpend {
6 pub session_id: String,
7 pub project: String,
8 pub latest_context_epoch: i64,
9 pub context_rows: i64,
10 pub context_output_chars: i64,
11 pub context_estimated_tokens: i64,
12 pub context_emit_count: i64,
13 pub context_suppress_count: i64,
14 pub relevance_state: String,
15 pub relevance_policy_version: Option<String>,
16 pub relevance_k: Option<i64>,
17 pub relevance_threshold: Option<f64>,
18 pub relevance_candidate_count: i64,
19 pub relevance_eligible_count: i64,
20 pub relevance_final_injected_count: i64,
21 pub relevance_below_threshold_count: i64,
22 pub relevance_k_limited_count: i64,
23 pub relevance_section_budget_count: i64,
24 pub relevance_total_char_limit_count: i64,
25 pub ai_usage_attribution: String,
26 pub ai_calls: i64,
27 pub ai_total_tokens: i64,
28 pub ai_estimated_cost_usd: f64,
29 pub ai_unattributed_legacy_calls: i64,
30}
31
32pub fn query_latest_session_memory_spend(
33 conn: &Connection,
34) -> Result<Option<LatestSessionMemorySpend>> {
35 if !crate::retrieval::temporal::sqlite_table_exists(conn, "context_injections")? {
36 return Ok(None);
37 }
38
39 let Some(session_id) = conn
40 .query_row(
41 "SELECT session_id
42 FROM context_injections
43 WHERE session_id IS NOT NULL
44 AND trim(session_id) <> ''
45 ORDER BY updated_at_epoch DESC, last_emitted_epoch DESC, id DESC
46 LIMIT 1",
47 [],
48 |row| row.get::<_, String>(0),
49 )
50 .optional()?
51 else {
52 return Ok(None);
53 };
54
55 let (
56 project,
57 latest_context_epoch,
58 context_rows,
59 context_output_chars,
60 context_emit_count,
61 context_suppress_count,
62 ) = conn.query_row(
63 "SELECT
64 (SELECT project
65 FROM context_injections latest
66 WHERE latest.session_id = ?1
67 ORDER BY latest.updated_at_epoch DESC,
68 latest.last_emitted_epoch DESC,
69 latest.id DESC
70 LIMIT 1),
71 COALESCE(MAX(updated_at_epoch), 0),
72 COUNT(*),
73 COALESCE(SUM(output_chars), 0),
74 COALESCE(SUM(emit_count), 0),
75 COALESCE(SUM(suppress_count), 0)
76 FROM context_injections
77 WHERE session_id = ?1",
78 params![session_id.as_str()],
79 |row| {
80 Ok((
81 row.get::<_, String>(0)?,
82 row.get::<_, i64>(1)?,
83 row.get::<_, i64>(2)?,
84 row.get::<_, i64>(3)?,
85 row.get::<_, i64>(4)?,
86 row.get::<_, i64>(5)?,
87 ))
88 },
89 )?;
90
91 let (
92 ai_usage_attribution,
93 ai_calls,
94 ai_total_tokens,
95 ai_estimated_cost_usd,
96 ai_unattributed_legacy_calls,
97 ) = if sqlite_column_exists(conn, "ai_usage_events", "session_id")? {
98 let (calls, total_tokens, estimated_cost_usd) = conn.query_row(
99 "SELECT COUNT(*),
100 COALESCE(SUM(total_tokens), 0),
101 COALESCE(SUM(estimated_cost_usd), 0.0)
102 FROM ai_usage_events
103 WHERE session_id = ?1",
104 params![session_id.as_str()],
105 |row| {
106 Ok((
107 row.get::<_, i64>(0)?,
108 row.get::<_, i64>(1)?,
109 row.get::<_, f64>(2)?,
110 ))
111 },
112 )?;
113 let unattributed_legacy_calls = conn.query_row(
114 "SELECT COUNT(*)
115 FROM ai_usage_events
116 WHERE session_id IS NULL",
117 [],
118 |row| row.get::<_, i64>(0),
119 )?;
120 let attribution = if unattributed_legacy_calls > 0 {
121 "partial"
122 } else {
123 "attributed"
124 };
125 (
126 attribution.to_string(),
127 calls,
128 total_tokens,
129 estimated_cost_usd,
130 unattributed_legacy_calls,
131 )
132 } else {
133 ("unavailable".to_string(), 0, 0, 0.0, 0)
134 };
135 let relevance = query_latest_relevance_spend(conn, &session_id)?;
136
137 Ok(Some(LatestSessionMemorySpend {
138 session_id,
139 project,
140 latest_context_epoch,
141 context_rows,
142 context_output_chars,
143 context_estimated_tokens: estimate_tokens_from_chars(context_output_chars),
144 context_emit_count,
145 context_suppress_count,
146 relevance_state: relevance.state,
147 relevance_policy_version: relevance.policy_version,
148 relevance_k: relevance.k,
149 relevance_threshold: relevance.threshold,
150 relevance_candidate_count: relevance.candidate_count,
151 relevance_eligible_count: relevance.eligible_count,
152 relevance_final_injected_count: relevance.final_injected_count,
153 relevance_below_threshold_count: relevance.below_threshold_count,
154 relevance_k_limited_count: relevance.k_limited_count,
155 relevance_section_budget_count: relevance.section_budget_count,
156 relevance_total_char_limit_count: relevance.total_char_limit_count,
157 ai_usage_attribution,
158 ai_calls,
159 ai_total_tokens,
160 ai_estimated_cost_usd,
161 ai_unattributed_legacy_calls,
162 }))
163}
164
165#[derive(Default)]
166struct LatestRelevanceSpend {
167 state: String,
168 policy_version: Option<String>,
169 k: Option<i64>,
170 threshold: Option<f64>,
171 candidate_count: i64,
172 eligible_count: i64,
173 final_injected_count: i64,
174 below_threshold_count: i64,
175 k_limited_count: i64,
176 section_budget_count: i64,
177 total_char_limit_count: i64,
178}
179
180fn query_latest_relevance_spend(
181 conn: &Connection,
182 session_id: &str,
183) -> Result<LatestRelevanceSpend> {
184 if !crate::retrieval::temporal::sqlite_table_exists(conn, "context_injection_items")? {
185 return Ok(LatestRelevanceSpend {
186 state: "unavailable".to_string(),
187 ..LatestRelevanceSpend::default()
188 });
189 }
190 let policy = conn
191 .query_row(
192 "SELECT injection_run_id, score, provenance
193 FROM context_injection_items
194 WHERE session_id = ?1
195 AND item_kind = 'sessionstart_relevance_policy'
196 ORDER BY injected_at_epoch DESC, id DESC
197 LIMIT 1",
198 params![session_id],
199 |row| {
200 Ok((
201 row.get::<_, String>(0)?,
202 row.get::<_, Option<f64>>(1)?,
203 row.get::<_, Option<String>>(2)?.unwrap_or_default(),
204 ))
205 },
206 )
207 .optional()?;
208 let Some((run_id, threshold, provenance)) = policy else {
209 return Ok(LatestRelevanceSpend {
210 state: "unavailable".to_string(),
211 ..LatestRelevanceSpend::default()
212 });
213 };
214 let values = provenance
215 .split(';')
216 .filter_map(|part| part.split_once('='))
217 .collect::<std::collections::HashMap<_, _>>();
218 let counts = conn.query_row(
219 "SELECT
220 COALESCE(SUM(CASE WHEN status = 'injected'
221 AND channel IN ('lessons', 'index', 'sessions')
222 THEN 1 ELSE 0 END), 0),
223 COALESCE(SUM(CASE WHEN drop_reason = 'below_sessionstart_relevance_threshold'
224 THEN 1 ELSE 0 END), 0),
225 COALESCE(SUM(CASE WHEN drop_reason = 'sessionstart_k_limit'
226 THEN 1 ELSE 0 END), 0),
227 COALESCE(SUM(CASE WHEN drop_reason = 'section_budget'
228 THEN 1 ELSE 0 END), 0),
229 COALESCE(SUM(CASE WHEN drop_reason = 'total_char_limit'
230 THEN 1 ELSE 0 END), 0)
231 FROM context_injection_items
232 WHERE injection_run_id = ?1",
233 params![run_id],
234 |row| {
235 Ok((
236 row.get::<_, i64>(0)?,
237 row.get::<_, i64>(1)?,
238 row.get::<_, i64>(2)?,
239 row.get::<_, i64>(3)?,
240 row.get::<_, i64>(4)?,
241 ))
242 },
243 )?;
244 Ok(LatestRelevanceSpend {
245 state: values
246 .get("state")
247 .copied()
248 .unwrap_or("unavailable")
249 .to_string(),
250 policy_version: values.get("policy").map(|value| (*value).to_string()),
251 k: values.get("k").and_then(|value| value.parse().ok()),
252 threshold,
253 candidate_count: parse_provenance_count(&values, "candidates"),
254 eligible_count: parse_provenance_count(&values, "eligible"),
255 final_injected_count: counts.0,
256 below_threshold_count: counts.1,
257 k_limited_count: counts.2,
258 section_budget_count: counts.3,
259 total_char_limit_count: counts.4,
260 })
261}
262
263fn parse_provenance_count(values: &std::collections::HashMap<&str, &str>, key: &str) -> i64 {
264 values
265 .get(key)
266 .and_then(|value| value.parse().ok())
267 .unwrap_or(0)
268}
269
270fn estimate_tokens_from_chars(chars: i64) -> i64 {
271 if chars <= 0 {
272 0
273 } else {
274 (chars + 3) / 4
275 }
276}
277
278fn sqlite_column_exists(conn: &Connection, table: &str, column: &str) -> Result<bool> {
279 if !crate::retrieval::temporal::sqlite_table_exists(conn, table)? {
280 return Ok(false);
281 }
282 conn.query_row(
283 "SELECT EXISTS(
284 SELECT 1 FROM pragma_table_info(?1)
285 WHERE name = ?2
286 )",
287 params![table, column],
288 |row| row.get(0),
289 )
290 .map_err(Into::into)
291}
292
293#[cfg(test)]
294mod tests {
295 use rusqlite::Connection;
296
297 use super::*;
298
299 fn setup_status_spend_schema(conn: &Connection, include_ai_session_id: bool) {
300 let ai_session_id_column = if include_ai_session_id {
301 "session_id TEXT,"
302 } else {
303 ""
304 };
305 conn.execute_batch(&format!(
306 "CREATE TABLE context_injections (
307 id INTEGER PRIMARY KEY,
308 host TEXT NOT NULL,
309 project TEXT NOT NULL,
310 injection_key TEXT NOT NULL,
311 session_id TEXT,
312 context_hash TEXT NOT NULL,
313 output_mode TEXT NOT NULL,
314 output_chars INTEGER NOT NULL,
315 created_at_epoch INTEGER NOT NULL,
316 updated_at_epoch INTEGER NOT NULL,
317 last_emitted_epoch INTEGER NOT NULL,
318 emit_count INTEGER NOT NULL DEFAULT 1,
319 suppress_count INTEGER NOT NULL DEFAULT 0
320 );
321 CREATE TABLE ai_usage_events (
322 id INTEGER PRIMARY KEY,
323 created_at TEXT NOT NULL,
324 created_at_epoch INTEGER NOT NULL,
325 project TEXT,
326 {ai_session_id_column}
327 operation TEXT NOT NULL,
328 executor TEXT NOT NULL,
329 model TEXT,
330 input_tokens INTEGER NOT NULL,
331 output_tokens INTEGER NOT NULL,
332 total_tokens INTEGER NOT NULL,
333 estimated_cost_usd REAL NOT NULL
334 );
335 CREATE TABLE context_injection_items (
336 id INTEGER PRIMARY KEY,
337 injection_run_id TEXT NOT NULL,
338 session_id TEXT,
339 item_kind TEXT NOT NULL,
340 channel TEXT NOT NULL,
341 score REAL,
342 status TEXT NOT NULL,
343 drop_reason TEXT,
344 provenance TEXT,
345 injected_at_epoch INTEGER NOT NULL
346 );"
347 ))
348 .expect("status spend schema should be created");
349 }
350
351 #[test]
352 fn latest_session_memory_spend_combines_context_and_ai_usage() -> Result<()> {
353 let conn = Connection::open_in_memory()?;
354 setup_status_spend_schema(&conn, true);
355 conn.execute_batch(
356 "INSERT INTO context_injections
357 (host, project, injection_key, session_id, context_hash, output_mode, output_chars,
358 created_at_epoch, updated_at_epoch, last_emitted_epoch, emit_count, suppress_count)
359 VALUES
360 ('codex-cli', '/old', 'old-key', 'sess-old', 'h1', 'full', 1200, 10, 10, 10, 1, 0),
361 ('codex-cli', '/repo', 'key-a', 'sess-new', 'h2', 'full', 801, 20, 31, 30, 2, 1),
362 ('codex-cli', '/repo', 'key-b', 'sess-new', 'h3', 'suppressed', 399, 21, 32, 29, 1, 3);
363 INSERT INTO ai_usage_events
364 (created_at, created_at_epoch, project, session_id, operation, executor, model,
365 input_tokens, output_tokens, total_tokens, estimated_cost_usd)
366 VALUES
367 ('2026-06-18T00:00:00Z', 30, '/repo', 'sess-new', 'summarize', 'codex-cli',
368 'codex-default', 100, 50, 150, 0.0015),
369 ('2026-06-18T00:00:01Z', 31, '/repo', 'sess-new', 'memory_candidate', 'codex-cli',
370 'codex-default', 60, 40, 100, 0.0010),
371 ('2026-06-18T00:00:02Z', 32, '/repo', NULL, 'legacy', 'codex-cli',
372 'codex-default', 999, 1, 1000, 9.0);",
373 )?;
374
375 let spend = query_latest_session_memory_spend(&conn)?
376 .ok_or_else(|| anyhow::anyhow!("latest session spend"))?;
377
378 assert_eq!(
379 spend,
380 LatestSessionMemorySpend {
381 session_id: "sess-new".to_string(),
382 project: "/repo".to_string(),
383 latest_context_epoch: 32,
384 context_rows: 2,
385 context_output_chars: 1200,
386 context_estimated_tokens: 300,
387 context_emit_count: 3,
388 context_suppress_count: 4,
389 relevance_state: "unavailable".to_string(),
390 relevance_policy_version: None,
391 relevance_k: None,
392 relevance_threshold: None,
393 relevance_candidate_count: 0,
394 relevance_eligible_count: 0,
395 relevance_final_injected_count: 0,
396 relevance_below_threshold_count: 0,
397 relevance_k_limited_count: 0,
398 relevance_section_budget_count: 0,
399 relevance_total_char_limit_count: 0,
400 ai_usage_attribution: "partial".to_string(),
401 ai_calls: 2,
402 ai_total_tokens: 250,
403 ai_estimated_cost_usd: 0.0025,
404 ai_unattributed_legacy_calls: 1,
405 }
406 );
407 Ok(())
408 }
409
410 #[test]
411 fn latest_session_memory_spend_uses_updated_activity_for_suppressed_sessions() -> Result<()> {
412 let conn = Connection::open_in_memory()?;
413 setup_status_spend_schema(&conn, true);
414 conn.execute_batch(
415 "INSERT INTO context_injections
416 (host, project, injection_key, session_id, context_hash, output_mode, output_chars,
417 created_at_epoch, updated_at_epoch, last_emitted_epoch, emit_count, suppress_count)
418 VALUES
419 ('codex-cli', '/old', 'old-key', 'sess-old', 'h1', 'full', 1200, 10, 100, 100, 1, 0),
420 ('codex-cli', '/repo', 'key-a', 'sess-new', 'h2', 'suppressed', 401, 20, 110, 90, 1, 2);",
421 )?;
422
423 let spend = query_latest_session_memory_spend(&conn)?
424 .ok_or_else(|| anyhow::anyhow!("latest session spend"))?;
425
426 assert_eq!(spend.session_id, "sess-new");
427 assert_eq!(spend.latest_context_epoch, 110);
428 assert_eq!(spend.context_suppress_count, 2);
429 Ok(())
430 }
431
432 #[test]
433 fn latest_session_memory_spend_reports_latest_relevance_policy_and_drops() -> Result<()> {
434 let conn = Connection::open_in_memory()?;
435 setup_status_spend_schema(&conn, true);
436 conn.execute_batch(
437 "INSERT INTO context_injections
438 (host, project, injection_key, session_id, context_hash, output_mode, output_chars,
439 created_at_epoch, updated_at_epoch, last_emitted_epoch, emit_count, suppress_count)
440 VALUES
441 ('codex-cli', '/repo', 'key-a', 'sess-new', 'h2', 'full', 401,
442 20, 31, 30, 1, 0);
443 INSERT INTO context_injection_items
444 (injection_run_id, session_id, item_kind, channel, score, status, drop_reason,
445 provenance, injected_at_epoch)
446 VALUES
447 ('run-1', 'sess-new', 'sessionstart_relevance_policy', 'policy', 0.5, 'injected',
448 NULL,
449 'policy=sessionstart_significant_token_v1;state=applied;k=1;threshold=0.500000;candidates=5;eligible=2;selected=1;below_threshold=3;k_limited=1',
450 31),
451 ('run-1', 'sess-new', 'memory', 'lessons', 0.8, 'injected', NULL, '', 31),
452 ('run-1', 'sess-new', 'memory', 'index', 0.7, 'dropped',
453 'sessionstart_k_limit', '', 31),
454 ('run-1', 'sess-new', 'session_summary', 'sessions', 0.0, 'dropped',
455 'below_sessionstart_relevance_threshold', '', 31);",
456 )?;
457
458 let spend = query_latest_session_memory_spend(&conn)?
459 .ok_or_else(|| anyhow::anyhow!("latest session spend"))?;
460
461 assert_eq!(spend.relevance_state, "applied");
462 assert_eq!(
463 spend.relevance_policy_version.as_deref(),
464 Some("sessionstart_significant_token_v1")
465 );
466 assert_eq!(spend.relevance_k, Some(1));
467 assert_eq!(spend.relevance_threshold, Some(0.5));
468 assert_eq!(spend.relevance_candidate_count, 5);
469 assert_eq!(spend.relevance_eligible_count, 2);
470 assert_eq!(spend.relevance_final_injected_count, 1);
471 assert_eq!(spend.relevance_below_threshold_count, 1);
472 assert_eq!(spend.relevance_k_limited_count, 1);
473 Ok(())
474 }
475
476 #[test]
477 fn latest_session_memory_spend_tolerates_legacy_ai_usage_without_session_column() -> Result<()>
478 {
479 let conn = Connection::open_in_memory()?;
480 setup_status_spend_schema(&conn, false);
481 conn.execute_batch(
482 "INSERT INTO context_injections
483 (host, project, injection_key, session_id, context_hash, output_mode, output_chars,
484 created_at_epoch, updated_at_epoch, last_emitted_epoch, emit_count, suppress_count)
485 VALUES
486 ('codex-cli', '/repo', 'key-a', 'sess-new', 'h2', 'full', 401, 20, 31, 30, 1, 0);",
487 )?;
488
489 let spend = query_latest_session_memory_spend(&conn)?
490 .ok_or_else(|| anyhow::anyhow!("latest session spend"))?;
491
492 assert_eq!(spend.ai_usage_attribution, "unavailable");
493 assert_eq!(spend.ai_calls, 0);
494 assert_eq!(spend.context_estimated_tokens, 101);
495 Ok(())
496 }
497
498 #[test]
499 fn latest_session_memory_spend_is_blank_without_context_rows() -> Result<()> {
500 let conn = Connection::open_in_memory()?;
501
502 assert!(query_latest_session_memory_spend(&conn)?.is_none());
503 Ok(())
504 }
505}