Skip to content
Sign up

flagon-io/g1t

public
0
355 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::links::{BRANCH_SESSIONS, FAMILY, FIRST_PROMPTS, INSERT_LINK, INSERT_OVERLAP, INSERT_TOUCH, TOUCHED_BY_OTHERS, end_tree_sql, tree_sql};
7use super::sessions::{INSERT_ENTRY, MOVE_ENTRIES, RENEW_CLAIM, harness_session_sql, insert_session_sql, refresh_sql, renewal, resume_sql};
8use super::{after_minutes, expire_claims_sql, time_out_sessions_sql};
9use g1t_contracts::time::rfc3339;
10use g1t_contracts::work::{claim_is_mine, session_overlaps, session_pair, subagent_harness_id};
11use std::collections::HashSet;
12
13const NOW: u64 = 1_791_800_000_000;
14
15/// Starts a subagent of `parent` the way `start_session` does, on issue 4.
16fn start_child(db: &Connection, id: &str, parent: &str, harness_id: &str, agent: &str) -> Option<String> {
17 let at = rfc3339(NOW);
18 db.query_row(
19 &insert_session_sql(),
20 params![
21 id, "rep_1", 4, None::<u32>, "main", "usr_1", "ana", "claude_code", harness_id, None::<String>, None::<String>, at, at,
22 parent, agent, "Find the endpoints", "Find every API endpoint and list them", None::<String>
23 ],
24 |row| row.get::<_, String>(0),
25 )
26 .optional()
27 .unwrap()
28}
29
30fn column(db: &Connection, sql: &str, params: impl rusqlite::Params) -> Vec<String> {
31 let mut statement = db.prepare(sql).unwrap();
32 statement.query_map(params, |row| row.get::<_, String>(0)).unwrap().map(Result::unwrap).collect()
33}
34
35fn touch(db: &Connection, session: &str, kind: &str, value: &str) -> bool {
36 db.query_row(INSERT_TOUCH, params![session, "rep_1", kind, value, rfc3339(NOW)], |_| Ok(())).optional().unwrap().is_some()
37}
38
39fn database() -> Connection {
40 let db = Connection::open_in_memory().unwrap();
41 db.execute_batch("PRAGMA foreign_keys = ON;").unwrap();
42 let dir = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
43 let mut names: Vec<String> = std::fs::read_dir(dir)
44 .unwrap()
45 .map(|entry| entry.unwrap().file_name().into_string().unwrap())
46 .filter(|name| name.ends_with(".sql"))
47 .collect();
48 names.sort();
49 for name in names {
50 let sql = std::fs::read_to_string(format!("{dir}/{name}")).unwrap();
51 db.execute_batch(&sql).unwrap_or_else(|error| panic!("{name}: {error}"));
52 }
53 db.execute(
54 "INSERT INTO pulls (id, repo_id, number, title, agent, runtime, author_id, author_name, created_at, updated_at)
55 VALUES ('pr_9', 'rep_1', 9, 'Greet', 'claude-code', 'external', 'usr_1', 'ana', '', '')",
56 [],
57 )
58 .unwrap();
59 db
60}
61
62fn start(db: &Connection, id: &str, harness_id: Option<&str>) -> Option<String> {
63 let at = rfc3339(NOW);
64 db.query_row(
65 &insert_session_sql(),
66 params![
67 id, "rep_1", 4, None::<u32>, "main", "usr_1", "ana", "claude_code", harness_id, None::<String>, None::<String>, at, at,
68 None::<String>, None::<String>, None::<String>, None::<String>, None::<String>
69 ],
70 |row| row.get::<_, String>(0),
71 )
72 .optional()
73 .unwrap()
74}
75
76/// Records an entry the way `record_session_entries` does: the session
77/// number it took, or none when its key was recorded before.
78fn record(db: &Connection, pull: Option<&str>, session: &str, text: &str, key: Option<&str>) -> Option<u32> {
79 db.query_row(INSERT_ENTRY, params![pull, "note", text, None::<String>, None::<String>, rfc3339(NOW), session, key], |row| {
80 row.get(0)
81 })
82 .optional()
83 .unwrap()
84}
85
86fn pull_entries(db: &Connection) -> Vec<(u32, String)> {
87 let mut rows = db.prepare("SELECT seq, text FROM session_entries WHERE pull_id = 'pr_9' ORDER BY seq").unwrap();
88 rows.query_map([], |row| Ok((row.get(0)?, row.get(1)?))).unwrap().map(Result::unwrap).collect()
89}
90
91#[test]
92fn starting_a_harness_session_again_returns_the_first() {
93 let db = database();
94 assert_eq!(start(&db, "ses_1", Some("claude-abc")).as_deref(), Some("ses_1"));
95 // The hook fired twice: nothing new, and the first is found by its id.
96 assert_eq!(start(&db, "ses_2", Some("claude-abc")), None);
97 let found: String = db
98 .query_row(&harness_session_sql(), params!["rep_1", "usr_1", "claude_code", "claude-abc"], |row| row.get(0))
99 .unwrap();
100 assert_eq!(found, "ses_1");
101 // Without a harness id every start is a session of its own.
102 assert_eq!(start(&db, "ses_3", None).as_deref(), Some("ses_3"));
103 assert_eq!(start(&db, "ses_4", None).as_deref(), Some("ses_4"));
104}
105
106#[test]
107fn a_quiet_session_times_out_and_a_heartbeat_brings_it_back() {
108 let db = database();
109 start(&db, "ses_1", Some("a"));
110 // Two hours and more without a heartbeat.
111 let late = rfc3339(NOW + g1t_contracts::work::SESSION_TIMEOUT_MS);
112 let ended: Vec<String> = db
113 .prepare(&time_out_sessions_sql())
114 .unwrap()
115 .query_map(params!["rep_1", late], |row| row.get(0))
116 .unwrap()
117 .map(Result::unwrap)
118 .collect();
119 assert_eq!(ended, ["ses_1"]);
120 let reason: String = db.query_row("SELECT end_reason FROM sessions WHERE id = 'ses_1'", [], |row| row.get(0)).unwrap();
121 assert_eq!(reason, "timed_out");
122 // Noticed once only.
123 assert_eq!(db.prepare(&time_out_sessions_sql()).unwrap().query_map(params!["rep_1", late], |_| Ok(())).unwrap().count(), 0);
124 // The harness reconnects.
125 let state: String = db
126 .query_row(
127 &resume_sql(),
128 params![
129 None::<u32>, 9, None::<String>, "abc123", None::<String>, None::<String>, None::<String>, None::<String>,
130 None::<String>, None::<String>, late, "ses_1"
131 ],
132 |row| row.get(10),
133 )
134 .unwrap();
135 assert_eq!(state, "active");
136}
137
138#[test]
139fn a_subagent_starts_once_and_once_finished_stays_finished() {
140 let db = database();
141 start(&db, "ses_parent", Some("claude-abc"));
142 let child = subagent_harness_id("claude-abc", "agent-def456");
143 assert_eq!(start_child(&db, "ses_c1", "ses_parent", &child, "Explore").as_deref(), Some("ses_c1"));
144 // SubagentStart, then the Agent tool's answer naming it without its prefix.
145 assert_eq!(start_child(&db, "ses_c2", "ses_parent", &subagent_harness_id("claude-abc", "def456"), "Explore"), None);
146 let found: String = db.query_row(&harness_session_sql(), params!["rep_1", "usr_1", "claude_code", child], |row| row.get(0)).unwrap();
147 assert_eq!(found, "ses_c1");
148 // SubagentStop ended it with its result.
149 db.execute("UPDATE sessions SET state = 'ended', end_reason = 'done', result = 'Found 12 endpoints' WHERE id = 'ses_c1'", [])
150 .unwrap();
151 // A late hook only fills in what it says: it stays ended, its result kept.
152 let (state, task, result): (String, String, String) = db
153 .query_row(
154 &refresh_sql(),
155 params![
156 None::<u32>, None::<u32>, None::<String>, None::<String>, None::<String>, None::<String>, "List endpoints",
157 None::<String>, None::<String>, None::<String>, rfc3339(NOW), "ses_c1"
158 ],
159 |row| Ok((row.get(10)?, row.get(18)?, row.get(20)?)),
160 )
161 .unwrap();
162 assert_eq!((state.as_str(), task.as_str(), result.as_str()), ("ended", "List endpoints", "Found 12 endpoints"));
163 // Its tree, from the parent.
164 assert_eq!(column(&db, &tree_sql(), params!["ses_parent"]), ["ses_c1"]);
165}
166
167#[test]
168fn ending_a_session_ends_the_subagents_it_left_going() {
169 let db = database();
170 start(&db, "ses_parent", Some("p"));
171 start_child(&db, "ses_c1", "ses_parent", "p/agent/1", "Explore");
172 start_child(&db, "ses_c2", "ses_parent", "p/agent/2", "Plan");
173 // A subagent of a subagent.
174 start_child(&db, "ses_c3", "ses_c1", "p/agent/3", "Explore");
175 db.execute("UPDATE sessions SET state = 'ended', end_reason = 'done' WHERE id = 'ses_c2'", []).unwrap();
176 let mut ended = column(&db, &end_tree_sql(), params!["ses_parent", rfc3339(NOW), "parent ended"]);
177 ended.sort();
178 assert_eq!(ended, ["ses_c1", "ses_c3"], "the one already done keeps its own end");
179 let reason: String = db.query_row("SELECT end_reason FROM sessions WHERE id = 'ses_c2'", [], |row| row.get(0)).unwrap();
180 assert_eq!(reason, "done");
181 let mut family = column(&db, FAMILY, params!["ses_c1"]);
182 family.sort();
183 assert_eq!(family, ["ses_c1", "ses_c3", "ses_parent"], "its parents up and its subagents down, not its sibling");
184}
185
186#[test]
187fn links_and_touches_are_kept_once_each() {
188 let db = database();
189 start(&db, "ses_1", Some("a"));
190 assert!(touch(&db, "ses_1", "file", "src/lib.rs"));
191 assert!(!touch(&db, "ses_1", "file", "src/lib.rs"), "said again: nothing new");
192 assert!(touch(&db, "ses_1", "branch", "4-greet"));
193 let link = |via: &str| db.execute(INSERT_LINK, params!["ses_1", "rep_1", "issue", 4, via, rfc3339(NOW)]).unwrap();
194 assert_eq!(link("claim"), 1);
195 assert_eq!(link("claim"), 0);
196 assert_eq!(link("ref"), 1);
197 // A pull request opened on the branch after the push finds the session.
198 db.execute("UPDATE pulls SET source_branch = '4-greet' WHERE id = 'pr_9'", []).unwrap();
199 assert_eq!(column(&db, BRANCH_SESSIONS, params!["rep_1", "[9]"]), ["ses_1"]);
200 assert!(column(&db, BRANCH_SESSIONS, params!["rep_1", "[10]"]).is_empty());
201}
202
203/// Twelve subagents of one session work on one issue, some on the same
204/// files: each is listed on the issue beside the others, each is told of
205/// every other that touched a file it touched (once per pair, for both
206/// sides), and each holds a claim of its own that the others see.
207#[test]
208fn twelve_subagents_on_one_issue_see_each_other() {
209 let db = database();
210 start(&db, "ses_parent", Some("lead"));
211 let children: Vec<String> = (1..=12).map(|n| format!("ses_c{n:02}")).collect();
212 for (n, id) in children.iter().enumerate() {
213 start_child(&db, id, "ses_parent", &subagent_harness_id("lead", &format!("agent-{n}")), "general-purpose").unwrap();
214 // Every one edits the shared router; each pair also shares a file.
215 touch(&db, id, "file", "src/router.rs");
216 touch(&db, id, "file", &format!("src/part{}.rs", n / 2));
217 touch(&db, id, "file", &format!("src/only{n}.rs"));
218 }
219 // All twelve, and their parent, are on the issue.
220 let on_issue = column(&db, "SELECT id FROM sessions WHERE repo_id = 'rep_1' AND issue = 4 ORDER BY id", []);
221 assert_eq!(on_issue.len(), 13);
222 let mut announced = 0;
223 for id in &children {
224 let n: usize = id[5..].parse().unwrap();
225 let mine = vec!["src/router.rs".to_owned(), format!("src/part{}.rs", (n - 1) / 2), format!("src/only{}.rs", n - 1)];
226 let touched: Vec<(String, String)> = {
227 let mut statement = db.prepare(TOUCHED_BY_OTHERS).unwrap();
228 statement
229 .query_map(params!["rep_1", serde_json::to_string(&mine).unwrap(), id], |row| Ok((row.get(0)?, row.get(1)?)))
230 .unwrap()
231 .map(Result::unwrap)
232 .collect()
233 };
234 let family: HashSet<String> = column(&db, FAMILY, params![id]).into_iter().collect();
235 let overlaps = session_overlaps(id, &family, &touched);
236 // Every sibling, and never its own parent.
237 assert_eq!(overlaps.len(), 11, "{id} sees its eleven siblings");
238 assert!(overlaps.iter().all(|(other, _)| other != "ses_parent" && other != id));
239 let partner = overlaps.iter().find(|(_, files)| files.len() == 2).expect("the sibling sharing a part file");
240 assert_eq!(partner.1, vec![format!("src/part{}.rs", (n - 1) / 2), "src/router.rs".to_owned()]);
241 for (other, files) in overlaps {
242 let (low, high) = session_pair(id, &other);
243 let new = db
244 .query_row(INSERT_OVERLAP, params!["rep_1", low, high, serde_json::to_string(&files).unwrap(), rfc3339(NOW)], |_| Ok(()))
245 .optional()
246 .unwrap()
247 .is_some();
248 // Each new pair is two overlap.detected events, one for each side.
249 if new {
250 announced += 2;
251 }
252 }
253 }
254 assert_eq!(announced, 12 * 11, "66 pairs, each announced to both sides once");
255 // Each claims the issue from its own session: the others see eleven.
256 let claims: Vec<(&str, Option<&str>)> = children.iter().map(|id| ("usr_1", Some(id.as_str()))).collect();
257 for id in &children {
258 let others = claims.iter().filter(|(actor, session)| !claim_is_mine(actor, *session, "usr_1", Some(id))).count();
259 assert_eq!(others, 11, "{id}");
260 }
261}
262
263#[test]
264fn entries_are_recorded_once_per_key_and_numbered_in_the_session() {
265 let db = database();
266 start(&db, "ses_1", Some("a"));
267 assert_eq!(record(&db, None, "ses_1", "planning", Some("k1")), Some(1));
268 // The hook retried: skipped.
269 assert_eq!(record(&db, None, "ses_1", "planning", Some("k1")), None);
270 assert_eq!(record(&db, None, "ses_1", "no key", None), Some(2));
271 assert_eq!(record(&db, None, "ses_1", "no key", None), Some(3));
272 // Another session numbers its own, and may reuse a key.
273 start(&db, "ses_2", Some("b"));
274 assert_eq!(record(&db, None, "ses_2", "other", Some("k1")), Some(1));
275}
276
277#[test]
278fn a_sessions_entries_join_its_pull_request_after_what_it_has() {
279 let db = database();
280 // Recorded on the pull request the old way.
281 db.execute(
282 "INSERT INTO session_entries (pull_id, seq, kind, text, at) VALUES ('pr_9', 1, 'note', 'legacy', '')",
283 [],
284 )
285 .unwrap();
286 start(&db, "ses_1", Some("a"));
287 record(&db, None, "ses_1", "before the pull request", Some("k1"));
288 record(&db, None, "ses_1", "still before", Some("k2"));
289 // The session names #9: what it recorded moves onto it, in order.
290 let base: u32 = db.query_row("SELECT COALESCE(MAX(seq), 0) FROM session_entries WHERE pull_id = 'pr_9'", [], |row| row.get(0)).unwrap();
291 db.execute(MOVE_ENTRIES, params!["pr_9", base, "ses_1"]).unwrap();
292 // Then it records on the pull request as it goes.
293 assert_eq!(record(&db, Some("pr_9"), "ses_1", "after", Some("k3")), Some(3));
294 assert_eq!(record(&db, Some("pr_9"), "ses_1", "after", Some("k3")), None);
295 assert_eq!(
296 pull_entries(&db),
297 [(1, "legacy".into()), (2, "before the pull request".into()), (3, "still before".into()), (4, "after".into())]
298 );
299}
300
301#[test]
302fn claims_expire_when_due_and_heartbeats_only_extend_them() {
303 let db = database();
304 let insert = |id: &str, expires_at: &str| {
305 db.execute(
306 "INSERT INTO claims (id, repo_id, issue, actor_id, actor_name, state, ttl_minutes, created_at, expires_at)
307 VALUES (?1, 'rep_1', 4, 'usr_1', 'ana', 'active', 30, '', ?2)",
308 params![id, expires_at],
309 )
310 .unwrap();
311 };
312 insert("clm_due", &rfc3339(NOW - 1));
313 insert("clm_live", &after_minutes(NOW, 10));
314 let expired: Vec<String> = db
315 .prepare(&expire_claims_sql())
316 .unwrap()
317 .query_map(params!["rep_1", rfc3339(NOW)], |row| row.get(0))
318 .unwrap()
319 .map(Result::unwrap)
320 .collect();
321 assert_eq!(expired, ["clm_due"]);
322 let (state, ended_at): (String, String) =
323 db.query_row("SELECT state, ended_at FROM claims WHERE id = 'clm_due'", [], |row| Ok((row.get(0)?, row.get(1)?))).unwrap();
324 assert_eq!((state.as_str(), ended_at), ("expired", rfc3339(NOW - 1)));
325 // A heartbeat renews the live one by its own length.
326 let renewed = renewal(NOW, 30, &after_minutes(NOW, 10)).unwrap();
327 assert_eq!(renewed, after_minutes(NOW, 30));
328 assert_eq!(db.execute(RENEW_CLAIM, params![renewed, "clm_live"]).unwrap(), 1);
329 // Never an expired one, and never shorter than it is.
330 assert_eq!(db.execute(RENEW_CLAIM, params![renewed, "clm_due"]).unwrap(), 0);
331 assert_eq!(renewal(NOW, 30, &after_minutes(NOW, 240)), None);
332}
333
334#[test]
335fn an_untitled_sessions_first_prompt_is_found_by_its_own_numbering() {
336 let db = database();
337 start(&db, "ses_1", Some("claude-1")).unwrap();
338 start(&db, "ses_2", Some("claude-2")).unwrap();
339 let entry = |session: &str, kind: &str, text: &str| {
340 db.query_row(INSERT_ENTRY, params![None::<String>, kind, text, None::<String>, None::<String>, rfc3339(NOW), session, None::<String>], |_| Ok(()))
341 .unwrap();
342 };
343 entry("ses_1", "tool_call", "Bash: ls");
344 entry("ses_1", "prompt", "Fix the preview builds that time out");
345 entry("ses_1", "prompt", "and the second one");
346 // ses_2 recorded no prompt.
347 entry("ses_2", "note", "a note");
348 let mut statement = db.prepare(FIRST_PROMPTS).unwrap();
349 let found: Vec<(String, String)> = statement
350 .query_map(params![r#"["ses_1","ses_2","ses_none"]"#, 12], |row| Ok((row.get(0)?, row.get(1)?)))
351 .unwrap()
352 .map(Result::unwrap)
353 .collect();
354 assert_eq!(found, vec![("ses_1".to_owned(), "Fix the prev".to_owned())]);
355}