| 1 | //! The statements coordination runs, against the schema every migration |
| 2 | //! leaves, in SQLite (which D1 is), with foreign keys on as D1 has them. |
| 3 | |
| 4 | use rusqlite::{Connection, OptionalExtension, params}; |
| 5 | |
| 6 | use super::sessions::{INSERT_ENTRY, MOVE_ENTRIES, RENEW_CLAIM, harness_session_sql, insert_session_sql, renewal, resume_sql}; |
| 7 | use super::{after_minutes, expire_claims_sql, time_out_sessions_sql}; |
| 8 | use g1t_contracts::time::rfc3339; |
| 9 | |
| 10 | const NOW: u64 = 1_791_800_000_000; |
| 11 | |
| 12 | fn database() -> Connection { |
| 13 | let db = Connection::open_in_memory().unwrap(); |
| 14 | db.execute_batch("PRAGMA foreign_keys = ON;").unwrap(); |
| 15 | let dir = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations"); |
| 16 | let mut names: Vec<String> = std::fs::read_dir(dir) |
| 17 | .unwrap() |
| 18 | .map(|entry| entry.unwrap().file_name().into_string().unwrap()) |
| 19 | .filter(|name| name.ends_with(".sql")) |
| 20 | .collect(); |
| 21 | names.sort(); |
| 22 | for name in names { |
| 23 | let sql = std::fs::read_to_string(format!("{dir}/{name}")).unwrap(); |
| 24 | db.execute_batch(&sql).unwrap_or_else(|error| panic!("{name}: {error}")); |
| 25 | } |
| 26 | db.execute( |
| 27 | "INSERT INTO pulls (id, repo_id, number, title, agent, runtime, author_id, author_name, created_at, updated_at) |
| 28 | VALUES ('pr_9', 'rep_1', 9, 'Greet', 'claude-code', 'external', 'usr_1', 'ana', '', '')", |
| 29 | [], |
| 30 | ) |
| 31 | .unwrap(); |
| 32 | db |
| 33 | } |
| 34 | |
| 35 | fn start(db: &Connection, id: &str, harness_id: Option<&str>) -> Option<String> { |
| 36 | let at = rfc3339(NOW); |
| 37 | db.query_row( |
| 38 | &insert_session_sql(), |
| 39 | params![id, "rep_1", 4, None::<u32>, "main", "usr_1", "ana", "claude_code", harness_id, None::<String>, None::<String>, at, at], |
| 40 | |row| row.get::<_, String>(0), |
| 41 | ) |
| 42 | .optional() |
| 43 | .unwrap() |
| 44 | } |
| 45 | |
| 46 | /// Records an entry the way `record_session_entries` does: the session |
| 47 | /// number it took, or none when its key was recorded before. |
| 48 | fn record(db: &Connection, pull: Option<&str>, session: &str, text: &str, key: Option<&str>) -> Option<u32> { |
| 49 | db.query_row(INSERT_ENTRY, params![pull, "note", text, None::<String>, None::<String>, rfc3339(NOW), session, key], |row| { |
| 50 | row.get(0) |
| 51 | }) |
| 52 | .optional() |
| 53 | .unwrap() |
| 54 | } |
| 55 | |
| 56 | fn pull_entries(db: &Connection) -> Vec<(u32, String)> { |
| 57 | let mut rows = db.prepare("SELECT seq, text FROM session_entries WHERE pull_id = 'pr_9' ORDER BY seq").unwrap(); |
| 58 | rows.query_map([], |row| Ok((row.get(0)?, row.get(1)?))).unwrap().map(Result::unwrap).collect() |
| 59 | } |
| 60 | |
| 61 | #[test] |
| 62 | fn starting_a_harness_session_again_returns_the_first() { |
| 63 | let db = database(); |
| 64 | assert_eq!(start(&db, "ses_1", Some("claude-abc")).as_deref(), Some("ses_1")); |
| 65 | // The hook fired twice: nothing new, and the first is found by its id. |
| 66 | assert_eq!(start(&db, "ses_2", Some("claude-abc")), None); |
| 67 | let found: String = db |
| 68 | .query_row(&harness_session_sql(), params!["rep_1", "usr_1", "claude_code", "claude-abc"], |row| row.get(0)) |
| 69 | .unwrap(); |
| 70 | assert_eq!(found, "ses_1"); |
| 71 | // Without a harness id every start is a session of its own. |
| 72 | assert_eq!(start(&db, "ses_3", None).as_deref(), Some("ses_3")); |
| 73 | assert_eq!(start(&db, "ses_4", None).as_deref(), Some("ses_4")); |
| 74 | } |
| 75 | |
| 76 | #[test] |
| 77 | fn a_quiet_session_times_out_and_a_heartbeat_brings_it_back() { |
| 78 | let db = database(); |
| 79 | start(&db, "ses_1", Some("a")); |
| 80 | // Two hours and more without a heartbeat. |
| 81 | let late = rfc3339(NOW + g1t_contracts::work::SESSION_TIMEOUT_MS); |
| 82 | let ended: Vec<String> = db |
| 83 | .prepare(&time_out_sessions_sql()) |
| 84 | .unwrap() |
| 85 | .query_map(params!["rep_1", late], |row| row.get(0)) |
| 86 | .unwrap() |
| 87 | .map(Result::unwrap) |
| 88 | .collect(); |
| 89 | assert_eq!(ended, ["ses_1"]); |
| 90 | let reason: String = db.query_row("SELECT end_reason FROM sessions WHERE id = 'ses_1'", [], |row| row.get(0)).unwrap(); |
| 91 | assert_eq!(reason, "timed_out"); |
| 92 | // Noticed once only. |
| 93 | assert_eq!(db.prepare(&time_out_sessions_sql()).unwrap().query_map(params!["rep_1", late], |_| Ok(())).unwrap().count(), 0); |
| 94 | // The harness reconnects. |
| 95 | let state: String = db |
| 96 | .query_row(&resume_sql(), params![None::<u32>, 9, None::<String>, "abc123", None::<String>, late, "ses_1"], |row| row.get(10)) |
| 97 | .unwrap(); |
| 98 | assert_eq!(state, "active"); |
| 99 | } |
| 100 | |
| 101 | #[test] |
| 102 | fn entries_are_recorded_once_per_key_and_numbered_in_the_session() { |
| 103 | let db = database(); |
| 104 | start(&db, "ses_1", Some("a")); |
| 105 | assert_eq!(record(&db, None, "ses_1", "planning", Some("k1")), Some(1)); |
| 106 | // The hook retried: skipped. |
| 107 | assert_eq!(record(&db, None, "ses_1", "planning", Some("k1")), None); |
| 108 | assert_eq!(record(&db, None, "ses_1", "no key", None), Some(2)); |
| 109 | assert_eq!(record(&db, None, "ses_1", "no key", None), Some(3)); |
| 110 | // Another session numbers its own, and may reuse a key. |
| 111 | start(&db, "ses_2", Some("b")); |
| 112 | assert_eq!(record(&db, None, "ses_2", "other", Some("k1")), Some(1)); |
| 113 | } |
| 114 | |
| 115 | #[test] |
| 116 | fn a_sessions_entries_join_its_pull_request_after_what_it_has() { |
| 117 | let db = database(); |
| 118 | // Recorded on the pull request the old way. |
| 119 | db.execute( |
| 120 | "INSERT INTO session_entries (pull_id, seq, kind, text, at) VALUES ('pr_9', 1, 'note', 'legacy', '')", |
| 121 | [], |
| 122 | ) |
| 123 | .unwrap(); |
| 124 | start(&db, "ses_1", Some("a")); |
| 125 | record(&db, None, "ses_1", "before the pull request", Some("k1")); |
| 126 | record(&db, None, "ses_1", "still before", Some("k2")); |
| 127 | // The session names #9: what it recorded moves onto it, in order. |
| 128 | let base: u32 = db.query_row("SELECT COALESCE(MAX(seq), 0) FROM session_entries WHERE pull_id = 'pr_9'", [], |row| row.get(0)).unwrap(); |
| 129 | db.execute(MOVE_ENTRIES, params!["pr_9", base, "ses_1"]).unwrap(); |
| 130 | // Then it records on the pull request as it goes. |
| 131 | assert_eq!(record(&db, Some("pr_9"), "ses_1", "after", Some("k3")), Some(3)); |
| 132 | assert_eq!(record(&db, Some("pr_9"), "ses_1", "after", Some("k3")), None); |
| 133 | assert_eq!( |
| 134 | pull_entries(&db), |
| 135 | [(1, "legacy".into()), (2, "before the pull request".into()), (3, "still before".into()), (4, "after".into())] |
| 136 | ); |
| 137 | } |
| 138 | |
| 139 | #[test] |
| 140 | fn claims_expire_when_due_and_heartbeats_only_extend_them() { |
| 141 | let db = database(); |
| 142 | let insert = |id: &str, expires_at: &str| { |
| 143 | db.execute( |
| 144 | "INSERT INTO claims (id, repo_id, issue, actor_id, actor_name, state, ttl_minutes, created_at, expires_at) |
| 145 | VALUES (?1, 'rep_1', 4, 'usr_1', 'ana', 'active', 30, '', ?2)", |
| 146 | params![id, expires_at], |
| 147 | ) |
| 148 | .unwrap(); |
| 149 | }; |
| 150 | insert("clm_due", &rfc3339(NOW - 1)); |
| 151 | insert("clm_live", &after_minutes(NOW, 10)); |
| 152 | let expired: Vec<String> = db |
| 153 | .prepare(&expire_claims_sql()) |
| 154 | .unwrap() |
| 155 | .query_map(params!["rep_1", rfc3339(NOW)], |row| row.get(0)) |
| 156 | .unwrap() |
| 157 | .map(Result::unwrap) |
| 158 | .collect(); |
| 159 | assert_eq!(expired, ["clm_due"]); |
| 160 | let (state, ended_at): (String, String) = |
| 161 | db.query_row("SELECT state, ended_at FROM claims WHERE id = 'clm_due'", [], |row| Ok((row.get(0)?, row.get(1)?))).unwrap(); |
| 162 | assert_eq!((state.as_str(), ended_at), ("expired", rfc3339(NOW - 1))); |
| 163 | // A heartbeat renews the live one by its own length. |
| 164 | let renewed = renewal(NOW, 30, &after_minutes(NOW, 10)).unwrap(); |
| 165 | assert_eq!(renewed, after_minutes(NOW, 30)); |
| 166 | assert_eq!(db.execute(RENEW_CLAIM, params![renewed, "clm_live"]).unwrap(), 1); |
| 167 | // Never an expired one, and never shorter than it is. |
| 168 | assert_eq!(db.execute(RENEW_CLAIM, params![renewed, "clm_due"]).unwrap(), 0); |
| 169 | assert_eq!(renewal(NOW, 30, &after_minutes(NOW, 240)), None); |
| 170 | } |