Skip to content
Sign up

flagon-io/g1t

public
0
337 linesCodeBlameRaw

Pick any line to see why it is the way it is: the commit, the pull request and issue it came from, and what the agent was thinking.

The runner base keeps its old comment, so the pushed base image still matches its folder1//! What the runner needs to apply a workspace's plan caps (see
2//! `@g1t/contracts` `compute.ts`): how many agents it has at work, what an
3//! issue's agents have spent, a comment from g1t when it cannot
4//! start, and somewhere for runs to wait for a free slot.
5//!
6//! The decisions are the runner's; this keeps the counts and the queue.
7
8use g1t_contracts::agents::*;
9use g1t_contracts::events::CommentCreated;
10use g1t_contracts::repos::RepoPath;
11use g1t_contracts::time::rfc3339;
12use g1t_contracts::{FailureCode, Outcome, User, new_id};
13use g1t_kit::now_ms;
14use serde::Deserialize;
15use worker::Result;
16use worker::wasm_bindgen::JsValue;
17
18use crate::Work;
19use crate::checks::hash;
20use crate::reviews::{AGENT_ID, AGENT_NAME};
21use crate::rows::{NumberRow, ValueRow};
22use crate::runs::member_of;
23
24/// The longest comment g1t leaves for a run it could not start.
25const MAX_COMMENT_CHARS: usize = 2000;
26/// The most runs one workspace may have waiting at once.
27const MAX_WAITING: u32 = 50;
28
29#[derive(Deserialize)]
30struct SpendRow {
31 spent: Option<f64>,
32}
33
34#[derive(Deserialize)]
35struct PullKey {
36 id: String,
37 issue_id: Option<String>,
38}
39
40#[derive(Deserialize)]
41struct WaitRow {
42 id: String,
43 workspace: String,
44 kind: String,
45 payload: String,
46 created_at: String,
47}
48
49#[derive(Deserialize)]
50struct CostRow {
51 cost_usd: Option<f64>,
52 token_hash: String,
53}
54
55/// The kinds that count against the agents-at-once cap, as SQL.
56fn agent_kinds_sql() -> String {
57 AGENT_KINDS
58 .iter()
59 .map(|kind| format!("'{}'", kind.as_str()))
60 .collect::<Vec<_>>()
61 .join(", ")
62}
63
64/// Dollars as millionths, rounded up.
65pub(crate) fn to_micros(usd: f64) -> i64 {
66 if usd.is_finite() && usd > 0.0 {
67 (usd * 1_000_000.0).ceil() as i64
68 } else {
69 0
70 }
71}
72
73impl Work {
74 async fn repo_id_of(&self, path: &RepoPath) -> Result<Outcome<String>> {
75 let service = User {
76 id: "svc_runner".to_owned(),
77 username: "g1t".to_owned(),
78 ..User::default()
79 };
80 Ok(match self.repo(path, &member_of(&service, &path.namespace)).await? {
81 Outcome::Ok(repo) => Outcome::Ok(repo.id),
82 Outcome::Fail(failure) => Outcome::Fail(failure),
83 })
84 }
85
86 pub(crate) async fn active_agents(&self, a: ActiveAgentsArgs) -> Result<u32> {
87 // A run whose sandbox died unnoticed does not hold a slot forever.
88 self.sweep_silent().await?;
89 let count = self
90 .db
91 .prepare(format!(
92 "SELECT count(*) AS n FROM agent_runs
93 WHERE workspace = ? AND status IN ('queued', 'running') AND kind IN ({})",
94 agent_kinds_sql()
95 ))
96 .bind(&[a.workspace.to_lowercase().into()])?
97 .first::<NumberRow>(None)
98 .await?;
99 Ok(count.map_or(0, |row| row.n))
100 }
101
102 pub(crate) async fn issue_spend(&self, a: IssueSpendArgs) -> Result<Outcome<IssueSpend>> {
103 let repo_id = match self.repo_id_of(&a.repo).await? {
104 Outcome::Ok(id) => id,
105 Outcome::Fail(failure) => return Ok(Outcome::Fail(failure)),
106 };
107 let pull = self
108 .db
109 .prepare("SELECT id, issue_id FROM pulls WHERE repo_id = ? AND number = ?")
110 .bind(&[repo_id.as_str().into(), a.number.into()])?
111 .first::<PullKey>(None)
112 .await?;
113 // The issue a pull request is for, or the pull request itself.
114 let issue_id = match &pull {
115 Some(pull) => pull.issue_id.clone(),
116 None => self
117 .db
118 .prepare("SELECT id AS value FROM issues WHERE repo_id = ? AND number = ?")
119 .bind(&[repo_id.as_str().into(), a.number.into()])?
120 .first::<ValueRow>(None)
121 .await?
122 .map(|row| row.value),
123 };
124 let (issue_number, spent) = match (&issue_id, &pull) {
125 (Some(issue_id), _) => {
126 let number = self
127 .db
128 .prepare("SELECT number AS n FROM issues WHERE id = ?")
129 .bind(&[issue_id.as_str().into()])?
130 .first::<NumberRow>(None)
131 .await?
132 .map_or(a.number, |row| row.n);
133 let spent = self
134 .db
135 .prepare(
136 "SELECT sum(cost_usd) AS spent FROM agent_runs
137 WHERE pull_id IN (SELECT id FROM pulls WHERE issue_id = ?1)
138 OR (repo_id = ?2 AND number = ?3 AND pull_id IS NULL)",
139 )
140 .bind(&[issue_id.as_str().into(), repo_id.as_str().into(), number.into()])?
141 .first::<SpendRow>(None)
142 .await?;
143 (number, spent)
144 }
145 (None, Some(pull)) => {
146 let spent = self
147 .db
148 .prepare(
149 "SELECT sum(cost_usd) AS spent FROM agent_runs
150 WHERE pull_id = ?1 OR (repo_id = ?2 AND number = ?3 AND pull_id IS NULL)",
151 )
152 .bind(&[pull.id.as_str().into(), repo_id.as_str().into(), a.number.into()])?
153 .first::<SpendRow>(None)
154 .await?;
155 (a.number, spent)
156 }
157 (None, None) => {
158 return Ok(Outcome::fail(FailureCode::NotFound, "No issue or pull request by that number."));
159 }
160 };
161 Ok(Outcome::Ok(IssueSpend {
162 issue: issue_number,
163 spent_micros: to_micros(spent.and_then(|row| row.spent).unwrap_or(0.0)),
164 }))
165 }
166
167 pub(crate) async fn wait_for_slot(&self, a: WaitForSlotArgs) -> Result<bool> {
168 // The claim is given back as if it had never been taken, so the
169 // revision it counted does not count.
170 let released = self
171 .db
172 .prepare(
173 "UPDATE pulls SET
174 revisions = CASE WHEN working_on = 'revision' AND revisions > 0 THEN revisions - 1 ELSE revisions END,
175 working_on = NULL, working_until = NULL, stage_detail = ?
176 WHERE id = ? AND status = 'open'
177 RETURNING id AS value",
178 )
179 .bind(&[a.reason.trim().into(), a.pull_id.as_str().into()])?
180 .first::<ValueRow>(None)
181 .await?;
182 Ok(released.is_some())
183 }
184
185 pub(crate) async fn agent_comment(&self, a: AgentCommentArgs) -> Result<bool> {
186 let body: String = a.body.trim().chars().take(MAX_COMMENT_CHARS).collect();
187 if body.is_empty() {
188 return Ok(false);
189 }
190 let repo_id = match self.repo_id_of(&a.repo).await? {
191 Outcome::Ok(id) => id,
192 Outcome::Fail(_) => return Ok(false),
193 };
194 let pull_id = self
195 .db
196 .prepare("SELECT id AS value FROM pulls WHERE repo_id = ? AND number = ?")
197 .bind(&[repo_id.as_str().into(), a.number.into()])?
198 .first::<ValueRow>(None)
199 .await?
200 .map(|row| row.value);
201 let table = if pull_id.is_some() { "pulls" } else { "issues" };
202 let now = now_ms();
203 let id = new_id("cmt", now);
204 let at = rfc3339(now);
205 self.db
206 .batch(vec![
207 self.db
208 .prepare(
209 "INSERT INTO comments (id, repo_id, number, author_id, author_name, body, created_at)
210 VALUES (?, ?, ?, ?, ?, ?, ?)",
211 )
212 .bind(&[
213 id.as_str().into(),
214 repo_id.as_str().into(),
215 a.number.into(),
216 AGENT_ID.into(),
217 AGENT_NAME.into(),
218 body.as_str().into(),
219 at.as_str().into(),
220 ])?,
221 self.db
222 .prepare(format!("UPDATE {table} SET updated_at = ? WHERE repo_id = ? AND number = ?"))
223 .bind(&[at.as_str().into(), repo_id.as_str().into(), a.number.into()])?,
224 ])
225 .await?;
226 self.publish_as(
227 "comment.created",
228 &repo_id,
229 Some(AGENT_ID.to_owned()),
230 CommentCreated {
231 comment_id: id,
232 repo_id: repo_id.clone(),
233 number: a.number,
234 pull_id,
235 verdict: None,
Cards you act on in chat; agents comment and review as themselves; names shown cleanly; commits on the calendar236 ..CommentCreated::default()
The runner base keeps its old comment, so the pushed base image still matches its folder237 },
238 )
239 .await?;
240 Ok(true)
241 }
242
243 pub(crate) async fn add_wait(&self, a: AddWaitArgs) -> Result<Outcome<bool>> {
244 let workspace = a.workspace.to_lowercase();
245 let waiting = self
246 .db
247 .prepare("SELECT count(*) AS n FROM agent_waits WHERE workspace = ?")
248 .bind(&[workspace.as_str().into()])?
249 .first::<NumberRow>(None)
250 .await?
251 .map_or(0, |row| row.n);
252 if waiting >= MAX_WAITING {
253 return Ok(Outcome::fail(
254 FailureCode::Conflict,
255 format!("{MAX_WAITING} runs are already waiting for this workspace's agents. Try again when some have started."),
256 ));
257 }
258 let now = now_ms();
259 self.db
260 .prepare("INSERT INTO agent_waits (id, workspace, kind, payload, created_at) VALUES (?, ?, ?, ?, ?)")
261 .bind(&[
262 new_id("wait", now).into(),
263 workspace.into(),
264 a.kind.into(),
265 serde_json::to_string(&a.payload)?.into(),
266 rfc3339(now).into(),
267 ])?
268 .run()
269 .await?;
270 Ok(Outcome::Ok(true))
271 }
272
273 pub(crate) async fn waiting_workspaces(&self, _a: WaitingWorkspacesArgs) -> Result<Vec<String>> {
274 let rows = self
275 .db
276 .prepare(
277 "SELECT workspace AS value FROM agent_waits
278 GROUP BY workspace ORDER BY min(created_at) LIMIT 100",
279 )
280 .all()
281 .await?
282 .results::<ValueRow>()?;
283 Ok(rows.into_iter().map(|row| row.value).collect())
284 }
285
286 pub(crate) async fn take_wait(&self, a: TakeWaitArgs) -> Result<Option<AgentWait>> {
287 let row = self
288 .db
289 .prepare(
290 "DELETE FROM agent_waits WHERE id = (
291 SELECT id FROM agent_waits WHERE workspace = ? ORDER BY created_at, id LIMIT 1)
292 RETURNING id, workspace, kind, payload, created_at",
293 )
294 .bind(&[a.workspace.to_lowercase().into()])?
295 .first::<WaitRow>(None)
296 .await?;
297 Ok(row.map(|row| AgentWait {
298 id: row.id,
299 workspace: row.workspace,
300 kind: row.kind,
301 payload: serde_json::from_str(&row.payload).unwrap_or(serde_json::Value::Null),
302 created_at: row.created_at,
303 }))
304 }
305
306 pub(crate) async fn run_cost(&self, a: RunCostArgs) -> Result<Option<f64>> {
307 let row = self
308 .db
309 .prepare("SELECT cost_usd, token_hash FROM agent_runs WHERE id = ?")
310 .bind(&[JsValue::from(a.run_id.as_str())])?
311 .first::<CostRow>(None)
312 .await?;
313 Ok(row
314 .filter(|row| row.token_hash == hash(&a.token))
315 .and_then(|row| row.cost_usd))
316 }
317}
318
319#[cfg(test)]
320mod tests {
321 use super::*;
322
323 #[test]
324 fn dollars_become_micros_rounded_up() {
325 assert_eq!(to_micros(0.094), 94_000);
326 assert_eq!(to_micros(0.0000001), 1);
327 assert_eq!(to_micros(-1.0), 0);
328 assert_eq!(to_micros(f64::NAN), 0);
329 }
330
331 #[test]
332 fn only_agent_kinds_hold_a_slot() {
333 let sql = agent_kinds_sql();
334 assert!(sql.contains("'implement'") && sql.contains("'plan'"));
335 assert!(!sql.contains("'checks'") && !sql.contains("'queue'") && !sql.contains("'mergecheck'"));
336 }
337}

This file's history is long; its oldest lines are credited to the oldest commit read.