oino_rpc/
backend.rs

1//! [`RpcBackend`]: a [`SessionBackend`] over an RPC connection, so remote frontends see the
2//! same commands and events as in-process ones.
3//!
4//! The wire drops the cumulative `message_update.message`/`partial`; the decoder rebuilds
5//! them from `message_start` and the deltas with [`PartialAssembler`] (the host's own
6//! assembler), so [`ClientEvent`]s are complete again.
7
8use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
9
10use futures::StreamExt;
11use futures::future::BoxFuture;
12use futures::stream::BoxStream;
13use oino_ai::stream::PartialAssembler;
14use oino_ai::{AgentMessage, JsonValue, ProviderEvent, Usage};
15use oino_client::protocol::ClientInfo;
16use oino_client::{ClientError, OinoCommand, RpcCommand, RpcResponseData, SessionBackend};
17use oino_core::events::{AgentSessionEvent, ClientEvent, OinoEvent};
18use oino_core::ui::UiResponse;
19
20use crate::client::{RpcClient, Transport};
21
22fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
23    mutex.lock().unwrap_or_else(PoisonError::into_inner)
24}
25
26/// Turns wire records back into [`ClientEvent`]s (one decoder per subscription).
27#[derive(Debug, Default)]
28pub struct EventDecoder {
29    assembler: Option<PartialAssembler>,
30    /// Decoders of nested agents (`agent_event`), by agent id.
31    agents: std::collections::HashMap<String, EventDecoder>,
32}
33
34impl EventDecoder {
35    /// Decodes one wire record; `None` for records that are not client events.
36    pub fn decode(&mut self, mut value: JsonValue) -> Option<ClientEvent> {
37        if value.get("type").and_then(JsonValue::as_str) == Some("agent_event") {
38            let agent_id = value.get("agentId")?.as_str()?.to_owned();
39            let inner = value.get_mut("event").map(JsonValue::take)?;
40            let event = self
41                .agents
42                .entry(agent_id.clone())
43                .or_default()
44                .decode(inner)?;
45            return Some(ClientEvent::Oino(OinoEvent::AgentEvent {
46                agent_id,
47                event: Box::new(event),
48            }));
49        }
50        if let Some(JsonValue::Array(agents)) = value
51            .get("type")
52            .and_then(JsonValue::as_str)
53            .filter(|t| *t == "agents_changed")
54            .and_then(|_| value.get("agents"))
55        {
56            let open: Vec<&str> = agents
57                .iter()
58                .filter_map(|a| a.get("id").and_then(JsonValue::as_str))
59                .collect();
60            self.agents.retain(|id, _| open.contains(&id.as_str()));
61        }
62        if value.get("type").and_then(JsonValue::as_str) == Some("message_update") {
63            let event: ProviderEvent =
64                serde_json::from_value(value.get("assistantMessageEvent")?.clone()).ok()?;
65            let assembler = self.assembler.as_mut()?;
66            let assistant_message_event = assembler.apply(event);
67            let mut partial = assembler.partial().clone();
68            if let Some(usage) = value
69                .get("usage")
70                .and_then(|u| serde_json::from_value::<Usage>(u.clone()).ok())
71            {
72                partial.usage = usage;
73            }
74            return Some(ClientEvent::Session(AgentSessionEvent::MessageUpdate {
75                message: AgentMessage::Assistant(partial),
76                assistant_message_event,
77            }));
78        }
79        let event: ClientEvent = serde_json::from_value(value).ok()?;
80        match &event {
81            ClientEvent::Session(AgentSessionEvent::MessageStart {
82                message: AgentMessage::Assistant(message),
83            }) => self.assembler = Some(PartialAssembler::new(message.clone())),
84            ClientEvent::Session(AgentSessionEvent::MessageEnd { .. }) => self.assembler = None,
85            _ => {}
86        }
87        Some(event)
88    }
89}
90
91/// A session reached over RPC (`oino --mode rpc` child or `oino serve`).
92#[derive(Clone)]
93pub struct RpcBackend {
94    client: RpcClient,
95    session_id: Arc<Mutex<String>>,
96}
97
98impl std::fmt::Debug for RpcBackend {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        f.debug_struct("RpcBackend")
101            .field("session_id", &*lock(&self.session_id))
102            .finish_non_exhaustive()
103    }
104}
105
106impl RpcBackend {
107    /// Connects and opts into Oino events with `_oino/hello`.
108    ///
109    /// # Errors
110    ///
111    /// Spawn or handshake failures.
112    pub async fn connect(
113        transport: Transport,
114        client: ClientInfo,
115        ui: Option<JsonValue>,
116    ) -> Result<Self, ClientError> {
117        Self::over(RpcClient::connect(transport)?, client, ui).await
118    }
119
120    /// Wraps a connected client (sends `_oino/hello`).
121    ///
122    /// # Errors
123    ///
124    /// Handshake failures.
125    pub async fn over(
126        rpc: RpcClient,
127        client: ClientInfo,
128        ui: Option<JsonValue>,
129    ) -> Result<Self, ClientError> {
130        let hello = rpc
131            .request(RpcCommand::Oino(OinoCommand::Hello {
132                client,
133                ui,
134                protocol: 1,
135            }))
136            .await?;
137        let session_id = hello
138            .as_ref()
139            .and_then(|h| h.get("sessionId"))
140            .and_then(JsonValue::as_str)
141            .ok_or_else(|| ClientError::Protocol("_oino/hello without sessionId".to_owned()))?
142            .to_owned();
143        let session_id = Arc::new(Mutex::new(session_id));
144        let tracked = Arc::downgrade(&session_id);
145        let mut events = rpc.events();
146        tokio::spawn(async move {
147            while let Some(event) = events.next().await {
148                if let Ok(ClientEvent::Oino(OinoEvent::SessionReplaced { session_id, .. })) =
149                    serde_json::from_value::<ClientEvent>(event)
150                {
151                    let Some(tracked) = tracked.upgrade() else {
152                        return;
153                    };
154                    *lock(&tracked) = session_id;
155                }
156            }
157        });
158        Ok(Self {
159            client: rpc,
160            session_id,
161        })
162    }
163
164    /// The typed client.
165    #[must_use]
166    pub const fn client(&self) -> &RpcClient {
167        &self.client
168    }
169}
170
171impl SessionBackend for RpcBackend {
172    fn send(&self, command: RpcCommand) -> BoxFuture<'_, Result<RpcResponseData, ClientError>> {
173        Box::pin(self.client.request(command))
174    }
175
176    fn events(&self) -> BoxStream<'static, ClientEvent> {
177        let mut decoder = EventDecoder::default();
178        self.client
179            .events()
180            .filter_map(move |value| futures::future::ready(decoder.decode(value)))
181            .boxed()
182    }
183
184    fn ui_respond(&self, response: UiResponse) {
185        let client = self.client.clone();
186        tokio::spawn(async move {
187            if let Err(error) = client.ui_respond(&response).await {
188                tracing::warn!("extension_ui_response failed: {error}");
189            }
190        });
191    }
192
193    fn session_id(&self) -> String {
194        lock(&self.session_id).clone()
195    }
196
197    fn shutdown(&self) -> BoxFuture<'_, ()> {
198        Box::pin(async move {
199            let _ = self.client.close().await;
200        })
201    }
202
203    /// Resolves on the server's `shutdown_requested` event.
204    fn shutdown_requested(&self) -> BoxFuture<'static, ()> {
205        let mut events = self.client.events();
206        Box::pin(async move {
207            while let Some(event) = events.next().await {
208                if event.get("type").and_then(JsonValue::as_str) == Some("shutdown_requested") {
209                    return;
210                }
211            }
212            std::future::pending::<()>().await;
213        })
214    }
215}
oino_rpc/backend.rs

oino_rpc/
backend.rs
