From 255a93c2e87875da2c7e92842a677133b52e4cad Mon Sep 17 00:00:00 2001
From: Bryan Helmkamp <bryan@brynary.com>
Date: Fri, 20 Feb 2026 12:57:00 -0400
Subject: [PATCH] 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>
---
 crates/unified-llm/src/lib.rs        |  3 +++
 crates/unified-llm/src/middleware.rs | 28 ++++++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/crates/unified-llm/src/lib.rs b/crates/unified-llm/src/lib.rs
index af92eeaf3..e56d1568a 100644
--- a/crates/unified-llm/src/lib.rs
+++ b/crates/unified-llm/src/lib.rs
@@ -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;
diff --git a/crates/unified-llm/src/middleware.rs b/crates/unified-llm/src/middleware.rs
index ea3d5fa02..2509f6b55 100644
--- a/crates/unified-llm/src/middleware.rs
+++ b/crates/unified-llm/src/middleware.rs
@@ -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)))
 }