fix: allow concurrent streamable http requests - #1186
Conversation
github-actions
Bot
added
T-documentation
nickcoai
changed the title
fac20e9 to
fd439b3
Compare
Keep cancellation and replies available while old session POSTs finish. Bound the wait for old POSTs and the replacement initialization handshake. Do not retry interrupted POSTs because the server may have processed them. Add regressions for recovery, queued cancellation, control timeouts, and server replies needed by active requests.
Match responses against all pending requests before removing a stream registration. Keep distinct numeric and string ids separate while preserving the existing fallback for servers that stringify numeric ids. Add a mixed-id regression and keep a separate registration owner alive in the abandoned-cancellation test.
| let control_generation = self.control_generation.load(Ordering::SeqCst); | ||
| let mut cancellation_target = None; | ||
| let registration = if W::supports_request_cancellation() { | ||
| match &item { | ||
| JsonRpcMessage::Request(request) => Some(RequestCancellationRegistration::new( | ||
| request.id.clone(), | ||
| self.ct.child_token(), | ||
| self.request_cancellations.clone(), | ||
| )), | ||
| JsonRpcMessage::Notification(notification) => { | ||
| cancellation_target = | ||
| self.cancel_request_from_notification(¬ification.notification); |
There was a problem hiding this comment.
A request from the old session may still be running when recovery moves to a new generation. If its handler finishes afterward, this code stamps the response with the new generation, and the client posts it to the replacement session.
oxcabe
left a comment
oxcabe
left a comment
There was a problem hiding this comment.
LGTM 👍🏼 added some non-blocking comments
| } | ||
| } | ||
|
|
||
| fn cancel_request_from_notification( |
There was a problem hiding this comment.
nb: this function (alongside is_control_message, supports_request_cancellation, and the control-generation counter) would make WorkerTransport less transport agnostic. That is, implying that the cancellation policy of one transport would now live inside the shared mechanisms.
What I'd suggest here is to keep WorkerTransport generic and put the interception in a streamable HTTP wrapper that implements Transport and delegates, e.g:
pub struct StreamableHttpClientTransport<C: StreamableHttpClient> {
inner: WorkerTransport<StreamableHttpClientWorker<C>>,
request_cancellations: RequestCancellations,
control_tx: mpsc::Sender<WorkerSendRequest<...>>,
control_generation: Arc<AtomicU64>,
}Now send() on the wrapper should perform the registration, the cancellation lookup, and the control-lane routing, then delegate to inner.
As a note, this change would require that StreamableHttpClientTransport becomes a NewType.
| } | ||
|
|
||
| /// Keep the same cancellation registration alive after the POST completes. | ||
| #[cfg(feature = "transport-streamable-http-client")] |
There was a problem hiding this comment.
nb: also a transport-specific feature as part of the generic worker model, similar situation to the cancel_request_from_notification comment down below.
| /// Advance the generation so the worker can reject older control sends. | ||
| pub fn advance_control_generation(&self) -> u64 { | ||
| self.control_generation | ||
| .fetch_add(1, Ordering::SeqCst) | ||
| .wrapping_add(1) | ||
| } |
There was a problem hiding this comment.
nb: similar situation with the rest of comments from this file, i.e. this might be too specific to the streamable HTTP implementation.
It's not that big of an issue if this stays here, but in that case I'd encourage expanding on the doc comment, explaining what the implications of "control generation" would be for any generic worker.
| type BoxedSseStream = BoxStream<'static, Result<Sse, SseError>>; | ||
| type SseTaskResult<E> = (Option<RequestId>, Result<(), StreamableHttpError<E>>); | ||
| const SESSION_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5); | ||
| const CONTROL_POST_TIMEOUT: Duration = Duration::from_secs(5); |
There was a problem hiding this comment.
nb: this should be made configurable, as a reply POST that runs longer than
five seconds would inevitably fail with ControlRequestTimeout.
My suggestion would be to approach it in a similar way to how session_recovery_timeout is defined in line 1976, promoting it into a config field:
/// Maximum time that a cancellation or reply POST can run after it starts
/// (default: five seconds).
pub control_request_timeout: Duration,
3 participants
Summary
Allow ordinary streamable http POSTs to run concurrently. Previously, one slow
application/jsonresponse blocked unrelated requests.StreamableHttpClientTransportConfig::max_concurrent_requestsdefaults to 16. Set it to1to keep ordinary POSTs serial;0is treated as1. Open sse response streams do not count against this limit. Initialization and protocol-version changes remain ordered with ordinary POSTs.Response cleanup matches ids against all pending requests, not just open response streams. Exact ids take precedence over the existing numeric-string fallback.
Cancellation and recovery
Cancellation and replies to server requests use a separate bounded queue with one extra POST slot. Each control POST has a five-second timeout after it starts. Queued and active ordinary POSTs can be cancelled without waiting for a free slot.
For an open legacy response stream, cancellation stops reading but keeps the stream alive until the cancellation send finishes or is dropped. This lets custom http adapters handle cancellation using their existing stream state. Ordinary legacy cancellation notifications still reach the adapter; modern http cancellation remains local.
When recovery is enabled, ordinary POSTs that report
SessionExpiredin the same session share one recovery attempt. Wait up tosession_recovery_timeoutfor old POSTs, then stop any that remain. Those interrupted POSTs returnSessionRecoveryTimeoutand are not retried: the server may have processed them. The replacement handshake has a separate timeout of the same length; both default to five seconds.Retry only ordinary POSTs that returned
SessionExpired, at most once each. Control POSTs are not retried. A request that already completed on the common response stream is not retried. Old queued controls cannot change replacement-session state, and a cancelled or dropped version-change request does not block later requests.The shared worker's control queue is opt-in; server scheduling is unchanged. Callers still decide which tools may run together and which need approval. Oauth refresh coordination remains a separate follow-up.
Validation
Passed locally:
cargo +1.96 test --offline --all-featuresjustfile: 35 tests covering concurrency, subscriptions, expired sessions, disconnects, connection reuse, and shutdown.cargo +1.96 clippy --offline --all-targets --all-features -- -D warningscargo +nightly fmt --all -- --checkgit diff --checkTwo downstream custom-adapter integration tests also passed without adapter source changes.