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-core/src/stream_session.rs
+37
View File
@@ -114,6 +114,9 @@ pub struct AsyncJobHandle {
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub expires_at: Option<OffsetDateTime>,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub last_poll_at: Option<OffsetDateTime>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
@@ -132,6 +135,15 @@ impl AsyncJobHandle {
matches!(self.status, JobStatus::Created | JobStatus::Running)
}
pub fn remaining_poll_delay_ms(&self, now: OffsetDateTime, poll_interval_ms: u64) -> u64 {
let Some(last_poll_at) = self.last_poll_at else {
return 0;
};
let elapsed_ms = (now - last_poll_at).whole_milliseconds().max(0) as u64;
poll_interval_ms.saturating_sub(elapsed_ms)
}
pub fn mark_finished(&mut self, now: OffsetDateTime, result: Value) {
self.status = JobStatus::Completed;
self.result = Some(result);
@@ -252,6 +264,7 @@ mod tests {
result: None,
error: None,
expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
last_poll_at: None,
created_at: timestamp("2026-04-06T12:00:00Z"),
updated_at: timestamp("2026-04-06T12:00:00Z"),
finished_at: None,
@@ -278,4 +291,28 @@ mod tests {
Some(timestamp("2026-04-06T12:02:00Z"))
);
}
#[test]
fn async_job_reports_remaining_poll_delay() {
let job = AsyncJobHandle {
id: AsyncJobId::new("job_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: OperationId::new("op_01"),
status: JobStatus::Running,
progress: json!({"percent": 60}),
result: None,
error: None,
expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
last_poll_at: Some(timestamp("2026-04-06T12:01:00Z")),
created_at: timestamp("2026-04-06T12:00:00Z"),
updated_at: timestamp("2026-04-06T12:00:00Z"),
finished_at: None,
};
assert_eq!(
job.remaining_poll_delay_ms(timestamp("2026-04-06T12:01:00.100Z"), 250),
150
);
}
}
crates/crank-registry/src/migrations.rs
+5
View File
@@ -558,6 +558,7 @@ pub async fn apply_postgres(pool: &PgPool) -> Result<(), sqlx::Error> {
result_json jsonb null,
error_json jsonb null,
expires_at timestamptz null,
last_poll_at timestamptz null,
created_at timestamptz not null,
updated_at timestamptz not null,
finished_at timestamptz null
@@ -566,6 +567,10 @@ pub async fn apply_postgres(pool: &PgPool) -> Result<(), sqlx::Error> {
.execute(pool)
.await?;
query("alter table async_jobs add column if not exists last_poll_at timestamptz null")
.execute(pool)
.await?;
query(
"create index if not exists async_jobs_workspace_status_idx
on async_jobs(workspace_id, status, updated_at desc)",
crates/crank-registry/src/postgres/mod.rs
+11
View File
@@ -556,6 +556,7 @@ fn map_async_job(row: &PgRow) -> Result<AsyncJobHandle, RegistryError> {
.try_get::<Option<Json<Value>>, _>("error_json")?
.map(|value| value.0),
expires_at: row.try_get("expires_at")?,
last_poll_at: row.try_get("last_poll_at")?,
created_at: row.try_get("created_at")?,
updated_at: row.try_get("updated_at")?,
finished_at: row.try_get("finished_at")?,
@@ -2171,6 +2172,15 @@ mod tests {
let loaded = registry.get_async_job(&job.id).await.unwrap().unwrap();
assert_eq!(loaded.status, JobStatus::Completed);
let touched = registry
.touch_async_job_poll(&job.id, &timestamp("2026-04-06T12:02:00Z"))
.await
.unwrap();
assert_eq!(
touched.last_poll_at,
Some(timestamp("2026-04-06T12:02:00Z"))
);
let page = registry
.list_async_jobs(AsyncJobFilter {
workspace_id: &test_workspace_id(),
@@ -2511,6 +2521,7 @@ mod tests {
result: None,
error: None,
expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
last_poll_at: None,
created_at: timestamp("2026-04-06T12:00:00Z"),
updated_at: timestamp("2026-04-06T12:00:00Z"),
finished_at: None,
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,