mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-06 05:12:03 +08:00
leookun
5cdf642dd1
- Updated the `append` function to include a flag for replacing closing requests, improving request handling. - Refactored the `run_sse_handler` and `bidi_handler` functions to utilize a new tracing mechanism, enhancing observability. - Introduced a new `trace_outcome` function to standardize tracing outcomes for requests. - Removed the `CursorTraceRecorder` in favor of a new `CursorTraceService` for better performance and non-blocking behavior. - Added tests to validate the new tracing functionality and ensure correct behavior during request processing.
54 lines
1.3 KiB
Rust
54 lines
1.3 KiB
Rust
//! Buffers, replays, broadcasts, and atomically closes downstream output.
|
|
|
|
use bytes::Bytes;
|
|
use tokio::sync::mpsc;
|
|
|
|
#[derive(Default)]
|
|
pub struct OutputHub {
|
|
state: parking_lot::Mutex<OutputState>,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct OutputState {
|
|
history: Vec<Bytes>,
|
|
subscribers: Vec<mpsc::UnboundedSender<Bytes>>,
|
|
closed: bool,
|
|
}
|
|
|
|
impl OutputHub {
|
|
pub fn emit(&self, frame: Bytes) -> bool {
|
|
let mut state = self.state.lock();
|
|
if state.closed {
|
|
return false;
|
|
}
|
|
state.history.push(frame.clone());
|
|
state
|
|
.subscribers
|
|
.retain(|subscriber| subscriber.send(frame.clone()).is_ok());
|
|
true
|
|
}
|
|
|
|
pub fn subscribe(&self) -> mpsc::UnboundedReceiver<Bytes> {
|
|
let (sender, receiver) = mpsc::unbounded_channel();
|
|
let mut state = self.state.lock();
|
|
for frame in &state.history {
|
|
let _ = sender.send(frame.clone());
|
|
}
|
|
if !state.closed {
|
|
state.subscribers.push(sender);
|
|
}
|
|
receiver
|
|
}
|
|
|
|
pub fn close(&self) -> bool {
|
|
let mut state = self.state.lock();
|
|
if state.closed {
|
|
return false;
|
|
}
|
|
state.closed = true;
|
|
state.subscribers.clear();
|
|
drop(state);
|
|
true
|
|
}
|
|
}
|