objectiveai_sdk/http/
notifier.rs

1//! Client-side handle for sending `client_request::Notify` frames
2//! over a streaming WS connection and awaiting their matching
3//! `client_response::Response` replies.
4//!
5//! Returned from every `*_streaming_ws` method alongside the chunk
6//! Stream. Drop the Notifier to relinquish the write half — when
7//! both the Notifier and the Stream are dropped, the demux task
8//! exits and the WS closes cleanly.
9
10use crate::client_objectiveai_mcp::{client_request, client_response};
11use futures::SinkExt;
12use futures::stream::SplitSink;
13use std::sync::Arc;
14use tokio::net::TcpStream;
15use tokio::sync::{Mutex, oneshot};
16use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, tungstenite};
17
18/// Shared sender half of a split WebSocket. Both the Notifier and
19/// the demux task (which writes `server_response` frames after
20/// handler dispatch) hold a clone; the mutex serializes writes so
21/// frames don't interleave on the wire.
22pub(crate) type SharedSink = Arc<
23    Mutex<
24        SplitSink<
25            WebSocketStream<MaybeTlsStream<TcpStream>>,
26            tungstenite::Message,
27        >,
28    >,
29>;
30
31/// Per-connection registry of outstanding notify ids and the
32/// oneshot senders the demux task fulfills when the matching
33/// `client_response::Response` arrives.
34pub(crate) type PendingNotifies =
35    Arc<dashmap::DashMap<String, oneshot::Sender<client_response::Response>>>;
36
37/// Client-side handle for sending notifies to a running agent
38/// completion (or any other streaming endpoint) over the same WS
39/// that carries the chunk stream.
40///
41/// Multiple in-flight notifies are supported; each gets a unique id
42/// and parks its own oneshot. Calls are independent and can run
43/// concurrently from multiple tasks.
44#[derive(Clone)]
45pub struct Notifier {
46    sink: SharedSink,
47    pending: PendingNotifies,
48}
49
50impl Notifier {
51    pub(crate) fn new(sink: SharedSink, pending: PendingNotifies) -> Self {
52        Self { sink, pending }
53    }
54
55    /// Forward a `notifications/{tools,resources}/list_changed`
56    /// observation from an upstream `mcp::Connection` up to the API,
57    /// which will fan it out as an SSE event on every matching
58    /// per-MCP GET stream — `/objectiveai` or
59    /// `/{owner}/{name}/{ver}/{mcp}`, subscribed under the
60    /// per-agent `response_id` + matching `McpKind`.
61    ///
62    /// Returns `Ok(())` if the server accepted the notify. Returns
63    /// the server-supplied `code + message` if the server replied
64    /// with an `Error` variant. Returns
65    /// [`super::HttpError::NotifyChannelClosed`] if the WS was torn
66    /// down before the reply arrived.
67    pub async fn notify_list_changed(
68        &self,
69        change: client_request::McpListChanged,
70    ) -> Result<(), super::HttpError> {
71        self.send(client_request::Payload::McpListChanged(change))
72            .await
73    }
74
75    /// Common send-and-await-ack body shared by every `notify_*` method.
76    async fn send(
77        &self,
78        payload: client_request::Payload,
79    ) -> Result<(), super::HttpError> {
80        let id = uuid::Uuid::new_v4().to_string();
81        let (tx, rx) = oneshot::channel();
82        self.pending.insert(id.clone(), tx);
83
84        let request = client_request::Request {
85            id: id.clone(),
86            payload,
87        };
88        let frame = match serde_json::to_string(&request) {
89            Ok(s) => s,
90            Err(e) => {
91                self.pending.remove(&id);
92                return Err(super::HttpError::NotifySerialize(e));
93            }
94        };
95
96        {
97            let mut guard = self.sink.lock().await;
98            if let Err(e) =
99                guard.send(tungstenite::Message::Text(frame.into())).await
100            {
101                drop(guard);
102                self.pending.remove(&id);
103                return Err(super::HttpError::NotifySend(e));
104            }
105        }
106
107        let response = match rx.await {
108            Ok(r) => r,
109            Err(_) => return Err(super::HttpError::NotifyChannelClosed),
110        };
111
112        match response {
113            client_response::Response::Ok { .. } => Ok(()),
114            client_response::Response::Error { code, message, .. } => {
115                Err(super::HttpError::NotifyRejected { code, message })
116            }
117        }
118    }
119}
objectiveai_sdk/http/notifier.rs

objectiveai_sdk/http/
notifier.rs
