bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mcp: throttle async job poll requests

Browse Source
This commit is contained in:
a.tolmachev
2026-05-01 16:48:02 +00:00
parent 4cd5d5b132
commit 66f28defe5
10 changed files with 312 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/crank-registry/src/postgres/stream.rs
+48 -4
View File
@@ -302,12 +302,13 @@ impl PostgresRegistry {
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at
) values (
$1, $2, $3, $4, $5, $6::jsonb, $7::jsonb, $8::jsonb, $9::timestamptz,
$10::timestamptz, $11::timestamptz, $12::timestamptz
$10::timestamptz, $11::timestamptz, $12::timestamptz, $13::timestamptz
)",
)
.bind(request.job.id.as_str())
@@ -319,6 +320,7 @@ impl PostgresRegistry {
.bind(request.job.result.clone().map(Json))
.bind(request.job.error.clone().map(Json))
.bind(request.job.expires_at)
.bind(request.job.last_poll_at)
.bind(request.job.created_at)
.bind(request.job.updated_at)
.bind(request.job.finished_at)
@@ -343,6 +345,7 @@ impl PostgresRegistry {
result_json,
error_json,
expires_at as \"expires_at: OffsetDateTime\",
last_poll_at as \"last_poll_at: OffsetDateTime\",
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\",
finished_at as \"finished_at: OffsetDateTime\"
@@ -364,6 +367,7 @@ impl PostgresRegistry {
result: row.result_json,
error: row.error_json,
expires_at: row.expires_at,
last_poll_at: row.last_poll_at,
created_at: row.created_at,
updated_at: row.updated_at,
finished_at: row.finished_at,
@@ -383,9 +387,9 @@ impl PostgresRegistry {
set status = $3,
progress_json = $4::jsonb,
result_json = $5::jsonb,
error_json = $6::jsonb,
expires_at = coalesce($7::timestamptz, expires_at),
updated_at = $8::timestamptz,
error_json = $6::jsonb,
expires_at = coalesce($7::timestamptz, expires_at),
updated_at = $8::timestamptz,
finished_at = case
when $9::timestamptz is null then finished_at
else $9::timestamptz
@@ -402,6 +406,7 @@ impl PostgresRegistry {
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at",
@@ -436,6 +441,43 @@ impl PostgresRegistry {
}
}
pub async fn touch_async_job_poll(
&self,
id: &AsyncJobId,
now: &OffsetDateTime,
) -> Result<AsyncJobHandle, RegistryError> {
let row = sqlx::query(
"update async_jobs
set last_poll_at = $2::timestamptz
where id = $1
returning
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at",
)
.bind(id.as_str())
.bind(now)
.fetch_optional(&self.pool)
.await?;
match row.as_ref() {
Some(row) => map_async_job(row),
None => Err(RegistryError::AsyncJobNotFound {
job_id: id.as_str().to_owned(),
}),
}
}
pub async fn cancel_async_job(
&self,
id: &AsyncJobId,
@@ -492,6 +534,7 @@ impl PostgresRegistry {
result_json,
error_json,
expires_at as \"expires_at: OffsetDateTime\",
last_poll_at as \"last_poll_at: OffsetDateTime\",
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\",
finished_at as \"finished_at: OffsetDateTime\"
@@ -524,6 +567,7 @@ impl PostgresRegistry {
result: row.result_json,
error: row.error_json,
expires_at: row.expires_at,
last_poll_at: row.last_poll_at,
created_at: row.created_at,
updated_at: row.updated_at,
finished_at: row.finished_at,