Sign in

pulls/pr_01m47d15m3e54sn21z27rpy5n9

public

pr_01m47d15m3e54sn21z27rpy5n9/services/work/src/compute.rs

336 lines12,049 bytesCodeBlame
1//! 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-agent 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-agent 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,
236 },
237 )
238 .await?;
239 Ok(true)
240 }
241
242 pub(crate) async fn add_wait(&self, a: AddWaitArgs) -> Result<Outcome<bool>> {
243 let workspace = a.workspace.to_lowercase();
244 let waiting = self
245 .db
246 .prepare("SELECT count(*) AS n FROM agent_waits WHERE workspace = ?")
247 .bind(&[workspace.as_str().into()])?
248 .first::<NumberRow>(None)
249 .await?
250 .map_or(0, |row| row.n);
251 if waiting >= MAX_WAITING {
252 return Ok(Outcome::fail(
253 FailureCode::Conflict,
254 format!("{MAX_WAITING} runs are already waiting for this workspace's agents. Try again when some have started."),
255 ));
256 }
257 let now = now_ms();
258 self.db
259 .prepare("INSERT INTO agent_waits (id, workspace, kind, payload, created_at) VALUES (?, ?, ?, ?, ?)")
260 .bind(&[
261 new_id("wait", now).into(),
262 workspace.into(),
263 a.kind.into(),
264 serde_json::to_string(&a.payload)?.into(),
265 rfc3339(now).into(),
266 ])?
267 .run()
268 .await?;
269 Ok(Outcome::Ok(true))
270 }
271
272 pub(crate) async fn waiting_workspaces(&self, _a: WaitingWorkspacesArgs) -> Result<Vec<String>> {
273 let rows = self
274 .db
275 .prepare(
276 "SELECT workspace AS value FROM agent_waits
277 GROUP BY workspace ORDER BY min(created_at) LIMIT 100",
278 )
279 .all()
280 .await?
281 .results::<ValueRow>()?;
282 Ok(rows.into_iter().map(|row| row.value).collect())
283 }
284
285 pub(crate) async fn take_wait(&self, a: TakeWaitArgs) -> Result<Option<AgentWait>> {
286 let row = self
287 .db
288 .prepare(
289 "DELETE FROM agent_waits WHERE id = (
290 SELECT id FROM agent_waits WHERE workspace = ? ORDER BY created_at, id LIMIT 1)
291 RETURNING id, workspace, kind, payload, created_at",
292 )
293 .bind(&[a.workspace.to_lowercase().into()])?
294 .first::<WaitRow>(None)
295 .await?;
296 Ok(row.map(|row| AgentWait {
297 id: row.id,
298 workspace: row.workspace,
299 kind: row.kind,
300 payload: serde_json::from_str(&row.payload).unwrap_or(serde_json::Value::Null),
301 created_at: row.created_at,
302 }))
303 }
304
305 pub(crate) async fn run_cost(&self, a: RunCostArgs) -> Result<Option<f64>> {
306 let row = self
307 .db
308 .prepare("SELECT cost_usd, token_hash FROM agent_runs WHERE id = ?")
309 .bind(&[JsValue::from(a.run_id.as_str())])?
310 .first::<CostRow>(None)
311 .await?;
312 Ok(row
313 .filter(|row| row.token_hash == hash(&a.token))
314 .and_then(|row| row.cost_usd))
315 }
316}
317
318#[cfg(test)]
319mod tests {
320 use super::*;
321
322 #[test]
323 fn dollars_become_micros_rounded_up() {
324 assert_eq!(to_micros(0.094), 94_000);
325 assert_eq!(to_micros(0.0000001), 1);
326 assert_eq!(to_micros(-1.0), 0);
327 assert_eq!(to_micros(f64::NAN), 0);
328 }
329
330 #[test]
331 fn only_agent_kinds_hold_a_slot() {
332 let sql = agent_kinds_sql();
333 assert!(sql.contains("'implement'") && sql.contains("'plan'"));
334 assert!(!sql.contains("'checks'") && !sql.contains("'queue'") && !sql.contains("'mergecheck'"));
335 }
336}