mcp: throttle rapid stream session polls
This commit is contained in:
@@ -58,6 +58,15 @@ impl StreamSession {
|
||||
!self.status.is_terminal() && !self.is_expired(now)
|
||||
}
|
||||
|
||||
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_polled(&mut self, now: OffsetDateTime) {
|
||||
self.last_poll_at = Some(now);
|
||||
}
|
||||
@@ -183,6 +192,54 @@ mod tests {
|
||||
assert!(!session.can_poll(timestamp("2026-04-06T12:03:00Z")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_poll_has_no_required_delay() {
|
||||
let session = StreamSession {
|
||||
id: StreamSessionId::new("stream_01"),
|
||||
workspace_id: WorkspaceId::new("ws_01"),
|
||||
agent_id: None,
|
||||
operation_id: OperationId::new("op_01"),
|
||||
protocol: Protocol::Rest,
|
||||
mode: ExecutionMode::Session,
|
||||
status: StreamStatus::Running,
|
||||
cursor: None,
|
||||
state: json!({"cursor":"abc"}),
|
||||
expires_at: timestamp("2026-04-06T12:05:00Z"),
|
||||
last_poll_at: None,
|
||||
created_at: timestamp("2026-04-06T12:00:00Z"),
|
||||
closed_at: None,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
session.remaining_poll_delay_ms(timestamp("2026-04-06T12:00:00.100Z"), 250),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rapid_repeat_poll_reports_remaining_delay() {
|
||||
let session = StreamSession {
|
||||
id: StreamSessionId::new("stream_01"),
|
||||
workspace_id: WorkspaceId::new("ws_01"),
|
||||
agent_id: None,
|
||||
operation_id: OperationId::new("op_01"),
|
||||
protocol: Protocol::Rest,
|
||||
mode: ExecutionMode::Session,
|
||||
status: StreamStatus::Running,
|
||||
cursor: None,
|
||||
state: json!({"cursor":"abc"}),
|
||||
expires_at: timestamp("2026-04-06T12:05:00Z"),
|
||||
last_poll_at: Some(timestamp("2026-04-06T12:01:00Z")),
|
||||
created_at: timestamp("2026-04-06T12:00:00Z"),
|
||||
closed_at: None,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
session.remaining_poll_delay_ms(timestamp("2026-04-06T12:01:00.100Z"), 250),
|
||||
150
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn async_job_tracks_finish_and_failure() {
|
||||
let mut job = AsyncJobHandle {
|
||||
|
||||
Reference in New Issue
Block a user