mcp: throttle async job poll requests

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
+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
);
}
}