Skip to content
Sign up

flagon-io/g1t

public
0
170 linesCodeBlameRaw
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
4use rusqlite::{Connection, OptionalExtension, params};
5
6use super::sessions::{INSERT_ENTRY, MOVE_ENTRIES, RENEW_CLAIM, harness_session_sql, insert_session_sql, renewal, resume_sql};
7use super::{after_minutes, expire_claims_sql, time_out_sessions_sql};
8use g1t_contracts::time::rfc3339;
9
10const NOW: u64 = 1_791_800_000_000;
11
12fn 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
35fn 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.
48fn 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
56fn 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]
62fn 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]
77fn 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]
102fn 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]
116fn 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]
140fn 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}