oino_rpc/
client.rs

1//! [`RpcClient`]: typed client for `oino --mode rpc`.
2
3use std::collections::HashMap;
4use std::path::PathBuf;
5use std::process::Stdio;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
8
9use futures::StreamExt;
10use futures::stream::BoxStream;
11use oino_ai::{AgentMessage, ImageContent, JsonValue, Model, ThinkingLevel};
12use oino_client::{
13    AgentCommand, ClientError, RpcCommand, RpcResponse, RpcResponseData, RpcSessionState,
14};
15use oino_core::compaction::CompactionResult;
16use oino_core::session::{
17    BashResult, CycleModelResult, ForkMessage, PromptDisposition, SessionStats, SlashCommandInfo,
18    StreamingBehavior,
19};
20use oino_core::ui::UiResponse;
21use serde::Deserialize;
22use serde::de::DeserializeOwned;
23use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
24use tokio::process::Child;
25use tokio::sync::{broadcast, oneshot};
26
27use crate::jsonl::{LineReader, serialize_line};
28
29/// Event buffer per subscriber.
30const EVENT_BUFFER: usize = 4096;
31
32fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
33    mutex.lock().unwrap_or_else(PoisonError::into_inner)
34}
35
36/// A child process speaking RPC on its stdio (`oino --mode rpc`).
37#[derive(Clone, Debug, Default, PartialEq, Eq)]
38pub struct ChildCommand {
39    /// Executable.
40    pub program: PathBuf,
41    /// Arguments (include `--mode rpc`).
42    pub args: Vec<String>,
43    /// Working directory.
44    pub cwd: Option<PathBuf>,
45    /// Extra environment.
46    pub env: Vec<(String, String)>,
47}
48
49/// Where the server is.
50#[derive(Clone, Debug, PartialEq, Eq)]
51#[non_exhaustive]
52pub enum Transport {
53    /// Spawn a child and talk over its stdin/stdout (stderr is inherited).
54    StdioChild(ChildCommand),
55}
56
57type Writer = Box<dyn AsyncWrite + Send + Unpin>;
58
59struct Inner {
60    /// `None` once closed (dropping the pipe is what signals EOF to a child).
61    writer: tokio::sync::Mutex<Option<Writer>>,
62    pending: Mutex<HashMap<String, oneshot::Sender<RpcResponse>>>,
63    next_id: AtomicU64,
64    events: Mutex<broadcast::Receiver<JsonValue>>,
65    child: Mutex<Option<Child>>,
66}
67
68/// Typed RPC client. Cheap to clone.
69#[derive(Clone)]
70pub struct RpcClient {
71    inner: Arc<Inner>,
72}
73
74impl std::fmt::Debug for RpcClient {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        f.debug_struct("RpcClient").finish_non_exhaustive()
77    }
78}
79
80/// `fork` data.
81#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
82pub struct ForkOutcome {
83    /// Text of the forked-from user message.
84    #[serde(default)]
85    pub text: Option<String>,
86    /// An extension cancelled.
87    pub cancelled: bool,
88}
89
90/// `clear_queue` data.
91#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
92#[serde(rename_all = "camelCase")]
93pub struct ClearedQueues {
94    /// Steering texts.
95    pub steering: Vec<String>,
96    /// Follow-up texts.
97    pub follow_up: Vec<String>,
98}
99
100/// `get_entries` / `get_tree` data.
101#[derive(Clone, Debug, PartialEq, Deserialize)]
102#[serde(rename_all = "camelCase")]
103pub struct EntriesOutcome {
104    /// Entries (`get_entries`) or tree nodes (`get_tree`), as JSON.
105    #[serde(alias = "tree")]
106    pub entries: Vec<JsonValue>,
107    /// Current leaf.
108    pub leaf_id: Option<String>,
109}
110
111#[derive(Deserialize)]
112struct Disposition {
113    disposition: PromptDisposition,
114}
115
116#[derive(Deserialize)]
117struct Cancelled {
118    cancelled: bool,
119}
120
121impl RpcClient {
122    /// Connects over `transport`.
123    ///
124    /// # Errors
125    ///
126    /// The child could not be spawned.
127    pub fn connect(transport: Transport) -> Result<Self, ClientError> {
128        match transport {
129            Transport::StdioChild(command) => {
130                let mut builder = tokio::process::Command::new(&command.program);
131                builder
132                    .args(&command.args)
133                    .envs(command.env.iter().map(|(k, v)| (k, v)))
134                    .stdin(Stdio::piped())
135                    .stdout(Stdio::piped())
136                    .stderr(Stdio::inherit())
137                    .kill_on_drop(true);
138                if let Some(cwd) = &command.cwd {
139                    builder.current_dir(cwd);
140                }
141                let mut child = builder
142                    .spawn()
143                    .map_err(|e| ClientError::Transport(e.to_string()))?;
144                let (Some(stdin), Some(stdout)) = (child.stdin.take(), child.stdout.take()) else {
145                    return Err(ClientError::Transport("child stdio unavailable".to_owned()));
146                };
147                let client = Self::from_io(stdout, stdin);
148                *lock(&client.inner.child) = Some(child);
149                Ok(client)
150            }
151        }
152    }
153
154    /// A client over an already connected byte stream pair.
155    pub fn from_io<R, W>(reader: R, writer: W) -> Self
156    where
157        R: AsyncRead + Unpin + Send + 'static,
158        W: AsyncWrite + Unpin + Send + 'static,
159    {
160        let (sender, receiver) = broadcast::channel(EVENT_BUFFER);
161        let inner = Arc::new(Inner {
162            writer: tokio::sync::Mutex::new(Some(Box::new(writer))),
163            pending: Mutex::new(HashMap::new()),
164            next_id: AtomicU64::new(1),
165            events: Mutex::new(receiver),
166            child: Mutex::new(None),
167        });
168        let weak = Arc::downgrade(&inner);
169        tokio::spawn(async move {
170            let mut lines = LineReader::new(reader);
171            while let Ok(Some(line)) = lines.next_line().await {
172                let Ok(value) = serde_json::from_str::<JsonValue>(&line) else {
173                    continue;
174                };
175                if let Some(response) = RpcResponse::from_value(&value) {
176                    let waiter = response
177                        .id
178                        .as_ref()
179                        .and_then(JsonValue::as_str)
180                        .and_then(|id| weak.upgrade().and_then(|i| lock(&i.pending).remove(id)));
181                    if let Some(waiter) = waiter {
182                        let _ = waiter.send(response);
183                        continue;
184                    }
185                }
186                let _ = sender.send(value);
187            }
188            if let Some(inner) = weak.upgrade() {
189                lock(&inner.pending).clear();
190            }
191        });
192        Self { inner }
193    }
194
195    /// Raw wire events (everything that is not a reply to this client), from now on.
196    pub fn events(&self) -> BoxStream<'static, JsonValue> {
197        let receiver = lock(&self.inner.events).resubscribe();
198        futures::stream::unfold(receiver, |mut receiver| async move {
199            loop {
200                match receiver.recv().await {
201                    Ok(event) => return Some((event, receiver)),
202                    Err(broadcast::error::RecvError::Lagged(skipped)) => {
203                        tracing::warn!(skipped, "rpc client lagged");
204                    }
205                    Err(broadcast::error::RecvError::Closed) => return None,
206                }
207            }
208        })
209        .boxed()
210    }
211
212    async fn write(&self, value: &JsonValue) -> Result<(), ClientError> {
213        let line = serialize_line(value);
214        let mut guard = self.inner.writer.lock().await;
215        let writer = guard.as_mut().ok_or(ClientError::Closed)?;
216        writer
217            .write_all(line.as_bytes())
218            .await
219            .map_err(|e| ClientError::Transport(e.to_string()))?;
220        writer
221            .flush()
222            .await
223            .map_err(|e| ClientError::Transport(e.to_string()))
224    }
225
226    /// Sends any command and waits for its reply.
227    ///
228    /// # Errors
229    ///
230    /// [`ClientError::Command`] with the server's error text, [`ClientError::Closed`] when
231    /// the server went away.
232    pub async fn request(&self, command: RpcCommand) -> Result<RpcResponseData, ClientError> {
233        let id = format!("req_{}", self.inner.next_id.fetch_add(1, Ordering::Relaxed));
234        let (tx, rx) = oneshot::channel();
235        lock(&self.inner.pending).insert(id.clone(), tx);
236        let request = oino_client::RpcRequest {
237            id: Some(JsonValue::String(id.clone())),
238            command,
239        };
240        if let Err(error) = self.write(&request.to_value()).await {
241            lock(&self.inner.pending).remove(&id);
242            return Err(error);
243        }
244        let response = rx.await.map_err(|_| ClientError::Closed)?;
245        response.result.map_err(ClientError::Command)
246    }
247
248    async fn typed<T: DeserializeOwned>(&self, command: AgentCommand) -> Result<T, ClientError> {
249        let data = self.request(command.into()).await?;
250        serde_json::from_value(data.unwrap_or(JsonValue::Null))
251            .map_err(|e| ClientError::Protocol(e.to_string()))
252    }
253
254    async fn unit(&self, command: AgentCommand) -> Result<(), ClientError> {
255        self.request(command.into()).await.map(|_| ())
256    }
257
258    /// Answers a dialog.
259    ///
260    /// # Errors
261    ///
262    /// The write failed.
263    pub async fn ui_respond(&self, response: &UiResponse) -> Result<(), ClientError> {
264        let value =
265            serde_json::to_value(response).map_err(|e| ClientError::Protocol(e.to_string()))?;
266        self.write(&value).await
267    }
268
269    /// Closes stdin (orderly server shutdown) and waits for the child to exit.
270    ///
271    /// # Errors
272    ///
273    /// Waiting for the child failed.
274    pub async fn close(&self) -> Result<Option<i32>, ClientError> {
275        if let Some(mut writer) = self.inner.writer.lock().await.take() {
276            let _ = writer.shutdown().await;
277        }
278        let child = lock(&self.inner.child).take();
279        match child {
280            Some(mut child) => child
281                .wait()
282                .await
283                .map(|status| status.code())
284                .map_err(|e| ClientError::Transport(e.to_string())),
285            None => Ok(None),
286        }
287    }
288
289    /// `prompt`.
290    ///
291    /// # Errors
292    ///
293    /// Prompt errors.
294    pub async fn prompt(
295        &self,
296        message: impl Into<String>,
297        images: Vec<ImageContent>,
298        streaming_behavior: Option<StreamingBehavior>,
299    ) -> Result<PromptDisposition, ClientError> {
300        let reply: Disposition = self
301            .typed(AgentCommand::Prompt {
302                message: message.into(),
303                images,
304                streaming_behavior,
305            })
306            .await?;
307        Ok(reply.disposition)
308    }
309
310    /// `steer`.
311    ///
312    /// # Errors
313    ///
314    /// Extension commands cannot be queued.
315    pub async fn steer(
316        &self,
317        message: impl Into<String>,
318        images: Vec<ImageContent>,
319    ) -> Result<PromptDisposition, ClientError> {
320        let reply: Disposition = self
321            .typed(AgentCommand::Steer {
322                message: message.into(),
323                images,
324            })
325            .await?;
326        Ok(reply.disposition)
327    }
328
329    /// `follow_up`.
330    ///
331    /// # Errors
332    ///
333    /// Extension commands cannot be queued.
334    pub async fn follow_up(
335        &self,
336        message: impl Into<String>,
337        images: Vec<ImageContent>,
338    ) -> Result<PromptDisposition, ClientError> {
339        let reply: Disposition = self
340            .typed(AgentCommand::FollowUp {
341                message: message.into(),
342                images,
343            })
344            .await?;
345        Ok(reply.disposition)
346    }
347
348    /// `abort`.
349    ///
350    /// # Errors
351    ///
352    /// Transport failures.
353    pub async fn abort(&self) -> Result<(), ClientError> {
354        self.unit(AgentCommand::Abort).await
355    }
356
357    /// `clear_queue`.
358    ///
359    /// # Errors
360    ///
361    /// Transport failures.
362    pub async fn clear_queue(&self) -> Result<ClearedQueues, ClientError> {
363        self.typed(AgentCommand::ClearQueue).await
364    }
365
366    /// `new_session`; `true` when an extension cancelled.
367    ///
368    /// # Errors
369    ///
370    /// Session file failures.
371    pub async fn new_session(&self, parent_session: Option<String>) -> Result<bool, ClientError> {
372        let reply: Cancelled = self
373            .typed(AgentCommand::NewSession { parent_session })
374            .await?;
375        Ok(reply.cancelled)
376    }
377
378    /// `get_state`.
379    ///
380    /// # Errors
381    ///
382    /// Transport failures.
383    pub async fn get_state(&self) -> Result<RpcSessionState, ClientError> {
384        self.typed(AgentCommand::GetState).await
385    }
386
387    /// `set_model`.
388    ///
389    /// # Errors
390    ///
391    /// `Model not found: …`, missing credentials.
392    pub async fn set_model(
393        &self,
394        provider: impl Into<String>,
395        model_id: impl Into<String>,
396    ) -> Result<Model, ClientError> {
397        self.typed(AgentCommand::SetModel {
398            provider: provider.into(),
399            model_id: model_id.into(),
400        })
401        .await
402    }
403
404    /// `cycle_model`.
405    ///
406    /// # Errors
407    ///
408    /// Persisting the selection failed.
409    pub async fn cycle_model(&self) -> Result<Option<CycleModelResult>, ClientError> {
410        self.typed(AgentCommand::CycleModel).await
411    }
412
413    /// `get_available_models`.
414    ///
415    /// # Errors
416    ///
417    /// Transport failures.
418    pub async fn get_available_models(&self) -> Result<Vec<Model>, ClientError> {
419        #[derive(Deserialize)]
420        struct Models {
421            models: Vec<Model>,
422        }
423        let reply: Models = self.typed(AgentCommand::GetAvailableModels).await?;
424        Ok(reply.models)
425    }
426
427    /// `set_thinking_level`.
428    ///
429    /// # Errors
430    ///
431    /// Transport failures.
432    pub async fn set_thinking_level(&self, level: ThinkingLevel) -> Result<(), ClientError> {
433        self.unit(AgentCommand::SetThinkingLevel { level }).await
434    }
435
436    /// `cycle_thinking_level`.
437    ///
438    /// # Errors
439    ///
440    /// Transport failures.
441    pub async fn cycle_thinking_level(&self) -> Result<Option<ThinkingLevel>, ClientError> {
442        #[derive(Deserialize)]
443        struct Level {
444            level: ThinkingLevel,
445        }
446        let reply: Option<Level> = self.typed(AgentCommand::CycleThinkingLevel).await?;
447        Ok(reply.map(|l| l.level))
448    }
449
450    /// `get_available_thinking_levels`.
451    ///
452    /// # Errors
453    ///
454    /// Transport failures.
455    pub async fn get_available_thinking_levels(&self) -> Result<Vec<ThinkingLevel>, ClientError> {
456        #[derive(Deserialize)]
457        struct Levels {
458            levels: Vec<ThinkingLevel>,
459        }
460        let reply: Levels = self.typed(AgentCommand::GetAvailableThinkingLevels).await?;
461        Ok(reply.levels)
462    }
463
464    /// `set_steering_mode`.
465    ///
466    /// # Errors
467    ///
468    /// Transport failures.
469    pub async fn set_steering_mode(&self, mode: oino_agent::QueueMode) -> Result<(), ClientError> {
470        self.unit(AgentCommand::SetSteeringMode { mode }).await
471    }
472
473    /// `set_follow_up_mode`.
474    ///
475    /// # Errors
476    ///
477    /// Transport failures.
478    pub async fn set_follow_up_mode(&self, mode: oino_agent::QueueMode) -> Result<(), ClientError> {
479        self.unit(AgentCommand::SetFollowUpMode { mode }).await
480    }
481
482    /// `compact`.
483    ///
484    /// # Errors
485    ///
486    /// Compaction errors.
487    pub async fn compact(
488        &self,
489        custom_instructions: Option<String>,
490    ) -> Result<CompactionResult, ClientError> {
491        self.typed(AgentCommand::Compact {
492            custom_instructions,
493        })
494        .await
495    }
496
497    /// `set_auto_compaction`.
498    ///
499    /// # Errors
500    ///
501    /// Transport failures.
502    pub async fn set_auto_compaction(&self, enabled: bool) -> Result<(), ClientError> {
503        self.unit(AgentCommand::SetAutoCompaction { enabled }).await
504    }
505
506    /// `set_auto_retry`.
507    ///
508    /// # Errors
509    ///
510    /// Transport failures.
511    pub async fn set_auto_retry(&self, enabled: bool) -> Result<(), ClientError> {
512        self.unit(AgentCommand::SetAutoRetry { enabled }).await
513    }
514
515    /// `abort_retry`.
516    ///
517    /// # Errors
518    ///
519    /// Transport failures.
520    pub async fn abort_retry(&self) -> Result<(), ClientError> {
521        self.unit(AgentCommand::AbortRetry).await
522    }
523
524    /// `bash`.
525    ///
526    /// # Errors
527    ///
528    /// The command could not run.
529    pub async fn bash(
530        &self,
531        command: impl Into<String>,
532        exclude_from_context: bool,
533    ) -> Result<BashResult, ClientError> {
534        self.typed(AgentCommand::Bash {
535            command: command.into(),
536            exclude_from_context: exclude_from_context.then_some(true),
537        })
538        .await
539    }
540
541    /// `abort_bash`.
542    ///
543    /// # Errors
544    ///
545    /// Transport failures.
546    pub async fn abort_bash(&self) -> Result<(), ClientError> {
547        self.unit(AgentCommand::AbortBash).await
548    }
549
550    /// `get_session_stats`.
551    ///
552    /// # Errors
553    ///
554    /// Transport failures.
555    pub async fn get_session_stats(&self) -> Result<SessionStats, ClientError> {
556        self.typed(AgentCommand::GetSessionStats).await
557    }
558
559    /// `export_html`; returns the written path.
560    ///
561    /// # Errors
562    ///
563    /// Write failures.
564    pub async fn export_html(&self, output_path: Option<String>) -> Result<String, ClientError> {
565        #[derive(Deserialize)]
566        struct Path {
567            path: String,
568        }
569        let reply: Path = self.typed(AgentCommand::ExportHtml { output_path }).await?;
570        Ok(reply.path)
571    }
572
573    /// `switch_session`; `true` when an extension cancelled.
574    ///
575    /// # Errors
576    ///
577    /// The file cannot be opened.
578    pub async fn switch_session(
579        &self,
580        session_path: impl Into<String>,
581    ) -> Result<bool, ClientError> {
582        let reply: Cancelled = self
583            .typed(AgentCommand::SwitchSession {
584                session_path: session_path.into(),
585            })
586            .await?;
587        Ok(reply.cancelled)
588    }
589
590    /// `fork`.
591    ///
592    /// # Errors
593    ///
594    /// `Invalid entry ID for forking`.
595    pub async fn fork(&self, entry_id: impl Into<String>) -> Result<ForkOutcome, ClientError> {
596        self.typed(AgentCommand::Fork {
597            entry_id: entry_id.into(),
598        })
599        .await
600    }
601
602    /// `clone`; `true` when an extension cancelled.
603    ///
604    /// # Errors
605    ///
606    /// `Cannot clone session: no current entry selected`.
607    pub async fn clone_session(&self) -> Result<bool, ClientError> {
608        let reply: Cancelled = self.typed(AgentCommand::Clone).await?;
609        Ok(reply.cancelled)
610    }
611
612    /// `get_fork_messages`.
613    ///
614    /// # Errors
615    ///
616    /// Transport failures.
617    pub async fn get_fork_messages(&self) -> Result<Vec<ForkMessage>, ClientError> {
618        #[derive(Deserialize)]
619        struct Messages {
620            messages: Vec<ForkMessage>,
621        }
622        let reply: Messages = self.typed(AgentCommand::GetForkMessages).await?;
623        Ok(reply.messages)
624    }
625
626    /// `get_entries`.
627    ///
628    /// # Errors
629    ///
630    /// `Entry not found: …`.
631    pub async fn get_entries(&self, since: Option<String>) -> Result<EntriesOutcome, ClientError> {
632        self.typed(AgentCommand::GetEntries { since }).await
633    }
634
635    /// `get_tree` (nodes as JSON in [`EntriesOutcome::entries`]).
636    ///
637    /// # Errors
638    ///
639    /// Transport failures.
640    pub async fn get_tree(&self) -> Result<EntriesOutcome, ClientError> {
641        self.typed(AgentCommand::GetTree).await
642    }
643
644    /// `get_last_assistant_text`.
645    ///
646    /// # Errors
647    ///
648    /// Transport failures.
649    pub async fn get_last_assistant_text(&self) -> Result<Option<String>, ClientError> {
650        #[derive(Deserialize)]
651        struct Text {
652            text: Option<String>,
653        }
654        let reply: Text = self.typed(AgentCommand::GetLastAssistantText).await?;
655        Ok(reply.text)
656    }
657
658    /// `set_session_name`.
659    ///
660    /// # Errors
661    ///
662    /// `Session name cannot be empty`.
663    pub async fn set_session_name(&self, name: impl Into<String>) -> Result<(), ClientError> {
664        self.unit(AgentCommand::SetSessionName { name: name.into() })
665            .await
666    }
667
668    /// `get_messages`.
669    ///
670    /// # Errors
671    ///
672    /// Transport failures.
673    pub async fn get_messages(&self) -> Result<Vec<AgentMessage>, ClientError> {
674        #[derive(Deserialize)]
675        struct Messages {
676            messages: Vec<AgentMessage>,
677        }
678        let reply: Messages = self.typed(AgentCommand::GetMessages).await?;
679        Ok(reply.messages)
680    }
681
682    /// `get_commands`.
683    ///
684    /// # Errors
685    ///
686    /// Transport failures.
687    pub async fn get_commands(&self) -> Result<Vec<SlashCommandInfo>, ClientError> {
688        #[derive(Deserialize)]
689        struct Commands {
690            commands: Vec<SlashCommandInfo>,
691        }
692        let reply: Commands = self.typed(AgentCommand::GetCommands).await?;
693        Ok(reply.commands)
694    }
695}
oino_rpc/client.rs

oino_rpc/
client.rs
