mcp: throttle async job poll requests
This commit is contained in:
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user