netget/client/msgpack_rpc/
mod.rs

1//! MessagePack-RPC client: one connection, pipelined calls matched by msgid, notifications in
2//! both directions.
3pub mod actions;
4use crate::client::llm_budget::call_llm_for_client;
5use crate::llm::actions::client_trait::{Client, ClientActionResult};
6use crate::logging::emit::Log;
7use crate::protocol::{ConnectContext, Event};
8use crate::server::msgpack_rpc::wire::{self, Rpc};
9use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
10use crate::state::{AccessLogOwner, ClientStatus};
11pub use actions::MsgpackRpcClientProtocol;
12use anyhow::{Context, Result};
13use serde_json::{json, Value};
14use std::collections::HashMap;
15use std::net::SocketAddr;
16use std::time::Duration;
17use tokio::{
18    io::{AsyncWriteExt, WriteHalf},
19    net::TcpStream,
20    sync::mpsc,
21};
22
23/// Calls awaiting a response at once.
24pub const MAX_PENDING: usize = 64;
25/// The connection may stay quiet indefinitely between messages.
26const READ_IDLE: Duration = Duration::from_secs(365 * 24 * 3600);
27
28pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
29    let stream = tokio::time::timeout(wire::IO_TIMEOUT, TcpStream::connect(&ctx.remote_addr))
30        .await
31        .context("MessagePack-RPC connect deadline")??;
32    let local = stream.local_addr()?;
33    let (reader, writer) = tokio::io::split(stream);
34    ctx.state
35        .update_client_status(ctx.client_id, ClientStatus::Connected)
36        .await;
37    let external =
38        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
39    let (msg_tx, msg_rx) = mpsc::channel::<Result<Value>>(64);
40    let reader_task = tokio::spawn(async move {
41        let mut stream = wire::Stream::new(reader);
42        loop {
43            let item = stream.next(READ_IDLE).await;
44            let end = !matches!(item, Ok(Some(_)));
45            let sent = match item {
46                Ok(Some(v)) => msg_tx.send(Ok(v)).await,
47                Ok(None) => {
48                    msg_tx
49                        .send(Err(anyhow::anyhow!("server closed the connection")))
50                        .await
51                }
52                Err(e) => msg_tx.send(Err(e)).await,
53            };
54            if end || sent.is_err() {
55                return;
56            }
57        }
58    });
59    ctx.state
60        .register_client_task(ctx.client_id, reader_task)
61        .await;
62    let (internal_tx, internal_rx) = mpsc::channel::<Value>(16);
63    let (event_tx, mut event_rx) = mpsc::channel::<Event>(64);
64    event_tx.try_send(Event::new(
65        &actions::CONNECTED_EVENT,
66        json!({"remote_addr": ctx.remote_addr}),
67    ))?;
68    let events_ctx = ctx.clone();
69    let dispatcher = tokio::spawn(async move {
70        let protocol = MsgpackRpcClientProtocol;
71        while let Some(event) = event_rx.recv().await {
72            let instruction = events_ctx
73                .state
74                .get_instruction_for_client(events_ctx.client_id)
75                .await
76                .unwrap_or_default();
77            let memory = events_ctx
78                .state
79                .get_memory_for_client(events_ctx.client_id)
80                .await
81                .unwrap_or_default();
82            match call_llm_for_client(
83                &events_ctx.llm_client,
84                &events_ctx.state,
85                events_ctx.client_id.to_string(),
86                &instruction,
87                &memory,
88                Some(&event),
89                &protocol,
90                &events_ctx.status_tx,
91            )
92            .await
93            {
94                Ok(result) => {
95                    if let Some(memory) = result.memory_updates {
96                        events_ctx
97                            .state
98                            .set_memory_for_client(events_ctx.client_id, memory)
99                            .await;
100                    }
101                    for action in result.actions {
102                        if internal_tx.send(action).await.is_err() {
103                            return;
104                        }
105                    }
106                }
107                Err(e) => Log::new(Some(&events_ctx.status_tx))
108                    .warn(format!("MessagePack-RPC client handler: {e}")),
109            }
110        }
111    });
112    let dispatcher_abort = dispatcher.abort_handle();
113    ctx.state
114        .register_client_task(ctx.client_id, dispatcher)
115        .await;
116    let session_ctx = ctx.clone();
117    let task = tokio::spawn(async move {
118        let result = session(
119            &session_ctx,
120            writer,
121            msg_rx,
122            external,
123            internal_rx,
124            event_tx,
125        )
126        .await;
127        if result.is_err() {
128            dispatcher_abort.abort();
129        }
130        let status = match result {
131            Ok(()) => ClientStatus::Disconnected,
132            Err(e) => {
133                Log::new(Some(&session_ctx.status_tx))
134                    .warn(format!("MessagePack-RPC client ended: {e}"));
135                ClientStatus::Error(e.to_string())
136            }
137        };
138        session_ctx
139            .state
140            .update_client_status(session_ctx.client_id, status)
141            .await;
142        session_ctx
143            .state
144            .remove_client_handle(session_ctx.client_id)
145            .await;
146        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
147    });
148    ctx.state.register_client_task(ctx.client_id, task).await;
149    Ok(local)
150}
151
152async fn session(
153    ctx: &ConnectContext,
154    mut writer: WriteHalf<TcpStream>,
155    mut messages: mpsc::Receiver<Result<Value>>,
156    mut external: mpsc::Receiver<ClientCommand>,
157    mut internal: mpsc::Receiver<Value>,
158    events: mpsc::Sender<Event>,
159) -> Result<()> {
160    let mut next_id: u64 = 0;
161    let mut pending: HashMap<u64, (String, Option<ClientCommand>)> = HashMap::new();
162    loop {
163        let (action, mut injected) = tokio::select! {
164            item = messages.recv() => {
165                let Some(item) = item else { return Ok(()) };
166                match wire::parse(item?)? {
167                    Rpc::Response { msgid, error, result } => {
168                        // A response to no call of ours is dropped.
169                        if let Some((method, caller)) = pending.remove(&msgid) {
170                            let event = json!({"method": method, "msgid": msgid, "result": result, "error": error});
171                            if let Some(command) = caller {
172                                crate::client::command_support::reply(command, Ok(ClientSendOutcome::Executed { detail: event.to_string() }));
173                            }
174                            events.try_send(Event::new(&actions::RESPONSE_EVENT, event)).context("MessagePack-RPC event queue full; consumer stalled")?;
175                        }
176                    }
177                    Rpc::Notification { method, params } => {
178                        events.try_send(Event::new(&actions::NOTIFICATION_EVENT, json!({"method": method, "params": params})))
179                            .context("MessagePack-RPC event queue full; consumer stalled")?;
180                    }
181                    Rpc::Request { msgid, method, .. } => {
182                        Log::new(Some(&ctx.status_tx)).warn(format!("MessagePack-RPC server called {method}; this client serves no methods"));
183                        writer.write_all(&wire::response(msgid, &json!("NetGet's client serves no methods"), &Value::Null)?).await?;
184                    }
185                }
186                continue;
187            }
188            command = external.recv() => match command {
189                Some(c) => (c.action.clone(), Some(c)),
190                None => return Ok(()),
191            },
192            action = internal.recv() => match action {
193                Some(a) => (a, None),
194                None => return Ok(()),
195            },
196        };
197        let refuse = |injected: &mut Option<ClientCommand>, error: String| {
198            if let Some(command) = injected.take() {
199                crate::client::command_support::reply(
200                    command,
201                    Ok(ClientSendOutcome::Rejected { error }),
202                );
203            }
204        };
205        match MsgpackRpcClientProtocol.execute_action(action.clone()) {
206            Ok(ClientActionResult::Disconnect) => {
207                if let Some(command) = injected.take() {
208                    crate::client::command_support::reply(
209                        command,
210                        Ok(ClientSendOutcome::Disconnected),
211                    );
212                }
213                return Ok(());
214            }
215            Ok(_) => {}
216            Err(e) => {
217                refuse(&mut injected, e.to_string());
218                continue;
219            }
220        }
221        if injected.is_some() {
222            ctx.state
223                .record_access_log(
224                    AccessLogOwner::Client(ctx.client_id.as_u32()),
225                    "MessagePack-RPC",
226                    None,
227                    "injected_action",
228                    action.clone(),
229                    vec![],
230                )
231                .await;
232        }
233        let method = action["method"].as_str().unwrap_or_default().to_string();
234        let bytes = if action["type"] == "msgpack_call" {
235            if pending.len() >= MAX_PENDING {
236                refuse(&mut injected, "too many calls awaiting responses".into());
237                continue;
238            }
239            next_id = (next_id + 1) % (u64::from(u32::MAX) + 1);
240            wire::request(next_id, &method, &action["params"])?
241        } else {
242            wire::notification(&method, &action["params"])?
243        };
244        tokio::time::timeout(wire::IO_TIMEOUT, writer.write_all(&bytes))
245            .await
246            .context("MessagePack-RPC write deadline")??;
247        if action["type"] == "msgpack_call" {
248            pending.insert(next_id, (method, injected.take()));
249        } else if let Some(command) = injected.take() {
250            crate::client::command_support::reply(
251                command,
252                Ok(ClientSendOutcome::Sent {
253                    bytes_sent: bytes.len(),
254                }),
255            );
256        }
257    }
258}
netget/client/msgpack_rpc/mod.rs

netget/client/msgpack_rpc/
mod.rs
