netget/server/msgpack_rpc/
mod.rs

1//! MessagePack-RPC server. Rust owns the MessagePack codec, the envelope, message ids and
2//! every bound; the handler answers each request and may send notifications back.
3pub mod actions;
4pub mod wire;
5
6use crate::llm::action_helper::call_llm;
7use crate::llm::actions::protocol_trait::ActionResult;
8use crate::logging::emit::Log;
9use crate::protocol::{Event, SpawnContext};
10use crate::server::connection::ConnectionId;
11use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
12use anyhow::{Context, Result};
13use serde_json::{json, Value};
14use std::{net::SocketAddr, time::Duration};
15use tokio::{
16    io::AsyncWriteExt,
17    net::{TcpListener, TcpStream},
18};
19use wire::Rpc;
20
21pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
22pub const MAX_IDLE_TIMEOUT_SECS: u64 = 86400;
23
24fn idle(ctx: &SpawnContext) -> Result<Duration> {
25    let secs = ctx
26        .startup_params
27        .as_ref()
28        .map(|p| p.get_optional_u64("idle_timeout_secs"))
29        .transpose()?
30        .flatten()
31        .unwrap_or(IDLE_TIMEOUT.as_secs());
32    anyhow::ensure!(
33        (1..=MAX_IDLE_TIMEOUT_SECS).contains(&secs),
34        "idle_timeout_secs must be between 1 and {MAX_IDLE_TIMEOUT_SECS}"
35    );
36    Ok(Duration::from_secs(secs))
37}
38
39pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
40    let idle = idle(&ctx)?;
41    let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
42    let addr = listener.local_addr()?;
43    Log::new(Some(&ctx.status_tx)).info(format!("MessagePack-RPC listening on {addr}"));
44    let state = ctx.state.clone();
45    let server_id = ctx.server_id;
46    let accept = tokio::spawn(async move {
47        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(
48            crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
49        );
50        loop {
51            let (socket, peer, permit) = match crate::server::accept_bounded::accept_bounded(
52                &listener,
53                &limiter,
54                b"",
55                "MessagePack-RPC",
56                Some(&ctx.status_tx),
57            )
58            .await
59            {
60                Ok(v) => v,
61                Err(_) => break,
62            };
63            let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
64            let now = crate::utils::clock::Instant::now();
65            ctx.state
66                .add_connection_to_server(
67                    server_id,
68                    ConnectionState {
69                        id,
70                        remote_addr: peer,
71                        local_addr: addr,
72                        bytes_sent: 0,
73                        bytes_received: 0,
74                        packets_sent: 0,
75                        packets_received: 0,
76                        last_activity: now,
77                        status: ConnectionStatus::Active,
78                        status_changed_at: now,
79                        protocol_info: ProtocolConnectionInfo::empty(),
80                    },
81                )
82                .await;
83            let child = ctx.clone();
84            ctx.state
85                .spawn_server_task(server_id, async move {
86                    let _permit = permit;
87                    if let Err(e) = session(&child, id, socket, peer, idle).await {
88                        Log::new(Some(&child.status_tx))
89                            .warn(format!("MessagePack-RPC connection {id} ended: {e}"));
90                    }
91                    child
92                        .state
93                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
94                        .await;
95                    let _ = child.status_tx.send("__UPDATE_UI__".into());
96                })
97                .await;
98        }
99    });
100    state.register_server_task(server_id, accept).await;
101    Ok(addr)
102}
103
104fn outcome(ctx: &SpawnContext, id: ConnectionId, operation: &str, decision: &str) {
105    let summary =
106        format!("MessagePack-RPC connection {id} operation={operation} decision={decision}");
107    let log = Log::new(Some(&ctx.status_tx));
108    if decision.starts_with("fail_closed_") || decision == "model_silent" {
109        log.error(summary);
110    } else {
111        log.info(summary);
112    }
113}
114
115/// What to write for one incoming message.
116async fn answer(
117    ctx: &SpawnContext,
118    id: ConnectionId,
119    rpc: Rpc,
120    peer: SocketAddr,
121) -> Result<Vec<u8>> {
122    let (event, msgid, op) = match &rpc {
123        Rpc::Request {
124            msgid,
125            method,
126            params,
127        } => (
128            Event::new(
129                &actions::REQUEST_EVENT,
130                json!({"method": method, "params": params, "msgid": msgid, "remote_addr": peer.to_string()}),
131            ),
132            Some(*msgid),
133            "msgpack_request",
134        ),
135        Rpc::Notification { method, params } => (
136            Event::new(
137                &actions::NOTIFICATION_EVENT,
138                json!({"method": method, "params": params, "remote_addr": peer.to_string()}),
139            ),
140            None,
141            "msgpack_notification",
142        ),
143        // A response from a client: NetGet's server sends no requests, so there is nothing
144        // it could answer.
145        Rpc::Response { .. } => anyhow::bail!("unsolicited response from the client"),
146    };
147    let failed = |error: Option<&anyhow::Error>| -> Result<Vec<u8>> {
148        match msgid {
149            Some(m) => {
150                let text = match error {
151                    Some(e) => crate::utils::wire_failure::prefixed_wire_failure_text(e),
152                    None => crate::utils::WireFailure::Unavailable.prefixed_text(),
153                };
154                wire::response(m, &json!(text), &Value::Null)
155            }
156            None => Ok(Vec::new()),
157        }
158    };
159    let result = match call_llm(
160        &ctx.llm_client,
161        &ctx.state,
162        ctx.server_id,
163        Some(id),
164        &event,
165        &actions::MsgpackRpcProtocol,
166    )
167    .await
168    {
169        Ok(r) => r,
170        Err(e) => {
171            outcome(ctx, id, op, "fail_closed_llm_error");
172            return failed(Some(&e));
173        }
174    };
175    if !result.failures.is_empty() {
176        outcome(ctx, id, op, "fail_closed_invalid_reply");
177        return failed(None);
178    }
179    let mut answers = Vec::new();
180    let mut notes = Vec::new();
181    let mut pending = result.protocol_results;
182    while let Some(r) = pending.pop() {
183        match r {
184            ActionResult::Custom { name, data } if name == "msgpack_notify" => notes.push(data),
185            ActionResult::Custom { name, data } if name.starts_with("msgpack_") => {
186                answers.push((name, data))
187            }
188            ActionResult::Multiple(items) => pending.extend(items),
189            _ => {}
190        }
191    }
192    notes.reverse();
193    let mut out = Vec::new();
194    match (msgid, answers.as_slice()) {
195        (Some(m), [(name, data)]) if name == "msgpack_result" => {
196            outcome(ctx, id, op, "model_answer");
197            out.extend(wire::response(m, &Value::Null, &data["result"])?);
198        }
199        (Some(m), [(name, data)]) if name == "msgpack_error" => {
200            outcome(ctx, id, op, "model_reject");
201            out.extend(wire::response(m, &data["error"], &Value::Null)?);
202        }
203        (Some(_), []) => {
204            outcome(ctx, id, op, "model_silent");
205            return failed(None);
206        }
207        (Some(_), _) => {
208            outcome(ctx, id, op, "fail_closed_invalid_reply");
209            return failed(None);
210        }
211        (None, []) => outcome(
212            ctx,
213            id,
214            op,
215            if notes.is_empty() {
216                "model_silent"
217            } else {
218                "model_answer"
219            },
220        ),
221        (None, [(name, _)]) if name == "msgpack_ignore" => outcome(ctx, id, op, "model_answer"),
222        (None, _) => {
223            outcome(ctx, id, op, "fail_closed_invalid_reply");
224            return Ok(Vec::new());
225        }
226    }
227    for n in notes {
228        out.extend(wire::notification(
229            n["method"].as_str().unwrap_or_default(),
230            &n["params"],
231        )?);
232    }
233    Ok(out)
234}
235
236async fn session(
237    ctx: &SpawnContext,
238    id: ConnectionId,
239    socket: TcpStream,
240    peer: SocketAddr,
241    idle: Duration,
242) -> Result<()> {
243    let (reader, mut writer) = tokio::io::split(socket);
244    let mut stream = wire::Stream::new(reader);
245    while let Some(value) = stream.next(idle).await? {
246        ctx.state
247            .update_connection_stats(ctx.server_id, id, None, None, Some(1), None)
248            .await;
249        let rpc = wire::parse(value)?;
250        let out = answer(ctx, id, rpc, peer).await?;
251        if !out.is_empty() {
252            tokio::time::timeout(wire::IO_TIMEOUT, writer.write_all(&out))
253                .await
254                .context("MessagePack-RPC write deadline")??;
255            ctx.state
256                .update_connection_stats(
257                    ctx.server_id,
258                    id,
259                    None,
260                    Some(out.len() as u64),
261                    None,
262                    Some(1),
263                )
264                .await;
265        }
266    }
267    Ok(())
268}
netget/server/msgpack_rpc/mod.rs

netget/server/msgpack_rpc/
mod.rs
