Add streaming middleware support and re-export set_default_client

Extend the Middleware trait with process_stream_event() for event-level
observation/transformation of streaming responses. Add
wrap_stream_with_middleware() helper. Re-export set_default_client at the
crate root per spec Section 2.5.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-02-20 12:57:00 -04:00
parent fa542605d7
commit 255a93c2e8
2 changed files with 29 additions and 2 deletions

View file

@ -10,3 +10,6 @@ pub mod catalog;
pub mod providers;
pub use tokio_util::sync::CancellationToken;
// Re-export module-level default client helpers (Section 2.5).
pub use generate::set_default_client;

View file

@ -1,6 +1,7 @@
use crate::error::SdkError;
use crate::provider::StreamEventStream;
use crate::types::{Request, Response};
use crate::types::{Request, Response, StreamEvent};
use futures::StreamExt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
@ -19,7 +20,11 @@ pub type NextStreamFn = Arc<
+ Sync,
>;
/// Middleware for intercepting `complete()` calls (Section 2.3).
/// Middleware for intercepting `complete()` and streaming calls (Section 2.3).
///
/// Implement `handle_complete` for blocking requests and `handle_stream` for
/// streaming requests. Override `process_stream` to observe or transform
/// individual stream events without replacing the entire stream handler.
#[async_trait::async_trait]
pub trait Middleware: Send + Sync {
async fn handle_complete(
@ -33,4 +38,23 @@ pub trait Middleware: Send + Sync {
request: Request,
next: NextStreamFn,
) -> Result<StreamEventStream, SdkError>;
/// Process an individual stream event. Override to observe or transform
/// events as they pass through the middleware. The default implementation
/// passes events through unchanged.
fn process_stream_event(
&self,
event: Result<StreamEvent, SdkError>,
) -> Result<StreamEvent, SdkError> {
event
}
}
/// Wrap a `StreamEventStream` so that each event passes through a middleware's
/// `process_stream_event` method.
pub fn wrap_stream_with_middleware(
stream: StreamEventStream,
middleware: Arc<dyn Middleware>,
) -> StreamEventStream {
Box::pin(stream.map(move |event| middleware.process_stream_event(event)))
}