netget/server/capnp_rpc/
mod.rs

1//! Cap'n Proto RPC server. Rust owns the encoding, the RPC tables and every bound; one
2//! capability (the bootstrap, typed by the startup schema's interface) is exported, and the
3//! handler answers each call on it with results or an exception.
4pub mod actions;
5pub mod layout;
6pub mod rpc;
7pub mod schema;
8
9use crate::llm::action_helper::call_llm;
10use crate::llm::actions::protocol_trait::ActionResult;
11use crate::logging::emit::Log;
12use crate::protocol::{Event, SpawnContext};
13use crate::server::connection::ConnectionId;
14use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
15use anyhow::{bail, ensure, Context, Result};
16use layout::{Message, Target};
17use rpc::{CallTarget, Incoming};
18use serde_json::{json, Value};
19use std::{collections::HashSet, net::SocketAddr, sync::Arc, time::Duration};
20use tokio::{
21    io::{AsyncWrite, AsyncWriteExt},
22    net::{TcpListener, TcpStream},
23};
24
25pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
26pub const MAX_IDLE_TIMEOUT_SECS: u64 = 86400;
27pub const WRITE_TIMEOUT: Duration = Duration::from_secs(10);
28/// How long `capnp compile` may take.
29pub const COMPILE_TIMEOUT: Duration = Duration::from_secs(30);
30/// Bootstrap questions one connection may hold open at once.
31pub const MAX_OPEN_BOOTSTRAPS: usize = 64;
32
33/// Longest inline schema accepted.
34pub const MAX_INLINE_SCHEMA_BYTES: usize = 256 * 1024;
35
36/// Schema source given inline rather than as a path: it has a declaration in it.
37fn is_inline(schema: &str) -> bool {
38    schema.contains('{') || schema.contains('\n')
39}
40
41/// Whether the configured schema is source that `capnp` must compile.
42pub fn needs_compiler(params: Option<&Value>) -> bool {
43    params
44        .and_then(|p| p.get("schema"))
45        .and_then(Value::as_str)
46        .is_none_or(|s| is_inline(s) || s.trim().ends_with(".capnp"))
47}
48
49async fn compile(path: &std::path::Path) -> Result<Vec<u8>> {
50    let mut command = tokio::process::Command::new("capnp");
51    command
52        .arg("compile")
53        .arg("-o-")
54        .arg(path)
55        .stdin(std::process::Stdio::null())
56        .kill_on_drop(true);
57    if let Some(dir) = path.parent().filter(|d| !d.as_os_str().is_empty()) {
58        command.arg(format!("--src-prefix={}", dir.display()));
59        command.arg(format!("--import-path={}", dir.display()));
60    }
61    let out = tokio::time::timeout(COMPILE_TIMEOUT, command.output())
62        .await
63        .context("capnp compile timed out")?
64        .context("could not run capnp (install capnproto)")?;
65    ensure!(
66        out.status.success(),
67        "capnp compile failed: {}",
68        crate::utils::truncate::truncate_for_log(String::from_utf8_lossy(&out.stderr).trim(), 1024)
69    );
70    Ok(out.stdout)
71}
72
73/// Load a schema: inline Cap'n Proto source, a `.capnp` file (both compiled with
74/// `capnp compile -o-`), or anything else as that command's output. Inline source without a
75/// file id gets one derived from its text, so the same text always has the same ids.
76pub async fn load_schema(schema: &str) -> Result<schema::Schema> {
77    let schema = schema.trim();
78    let bytes = if is_inline(schema) {
79        ensure!(
80            schema.len() <= MAX_INLINE_SCHEMA_BYTES,
81            "inline schema longer than {MAX_INLINE_SCHEMA_BYTES} bytes"
82        );
83        let has_id = schema
84            .lines()
85            .map(str::trim)
86            .any(|l| l.starts_with("@0x") && l.ends_with(';'));
87        let source = if has_id {
88            schema.to_string()
89        } else {
90            use std::hash::{Hash, Hasher};
91            let mut h = std::collections::hash_map::DefaultHasher::new();
92            schema.hash(&mut h);
93            format!("@{:#018x};\n{schema}\n", h.finish() | 1 << 63)
94        };
95        let dir = tempfile::tempdir()?;
96        let file = dir.path().join("inline.capnp");
97        tokio::fs::write(&file, source).await?;
98        compile(&file).await?
99    } else if schema.ends_with(".capnp") {
100        compile(std::path::Path::new(schema)).await?
101    } else {
102        let meta = tokio::fs::metadata(schema)
103            .await
104            .with_context(|| format!("cannot read schema {schema}"))?;
105        ensure!(
106            meta.len() <= schema::MAX_SCHEMA_BYTES as u64,
107            "compiled schema too large"
108        );
109        tokio::fs::read(schema).await?
110    };
111    schema::Schema::load(&bytes)
112}
113
114pub struct Config {
115    pub schema: schema::Schema,
116    pub interface: u64,
117    pub interface_name: String,
118    pub idle: Duration,
119}
120
121async fn config(ctx: &SpawnContext) -> Result<Config> {
122    let params = ctx
123        .startup_params
124        .as_ref()
125        .context("schema and interface are required")?;
126    let path = params.get_string("schema")?;
127    let name = params.get_string("interface")?;
128    let secs = params
129        .get_optional_u64("idle_timeout_secs")?
130        .unwrap_or(IDLE_TIMEOUT.as_secs());
131    ensure!(
132        (1..=MAX_IDLE_TIMEOUT_SECS).contains(&secs),
133        "idle_timeout_secs must be between 1 and {MAX_IDLE_TIMEOUT_SECS}"
134    );
135    let schema = load_schema(&path).await?;
136    let iface = schema.interface(&name)?;
137    let (interface, interface_name) = (iface.id, iface.name.clone());
138    Ok(Config {
139        schema,
140        interface,
141        interface_name,
142        idle: Duration::from_secs(secs),
143    })
144}
145
146pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
147    let config = Arc::new(config(&ctx).await?);
148    let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
149    let addr = listener.local_addr()?;
150    Log::new(Some(&ctx.status_tx)).info(format!(
151        "Cap'n Proto RPC listening on {addr}, exporting {}",
152        config.interface_name
153    ));
154    let state = ctx.state.clone();
155    let server_id = ctx.server_id;
156    let accept = tokio::spawn(async move {
157        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(
158            crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
159        );
160        loop {
161            let (socket, peer, permit) = match crate::server::accept_bounded::accept_bounded(
162                &listener,
163                &limiter,
164                b"",
165                "Cap'n Proto RPC",
166                Some(&ctx.status_tx),
167            )
168            .await
169            {
170                Ok(v) => v,
171                Err(_) => break,
172            };
173            let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
174            let now = crate::utils::clock::Instant::now();
175            ctx.state
176                .add_connection_to_server(
177                    server_id,
178                    ConnectionState {
179                        id,
180                        remote_addr: peer,
181                        local_addr: addr,
182                        bytes_sent: 0,
183                        bytes_received: 0,
184                        packets_sent: 0,
185                        packets_received: 0,
186                        last_activity: now,
187                        status: ConnectionStatus::Active,
188                        status_changed_at: now,
189                        protocol_info: ProtocolConnectionInfo::empty(),
190                    },
191                )
192                .await;
193            let child = ctx.clone();
194            let config = config.clone();
195            ctx.state
196                .spawn_server_task(server_id, async move {
197                    let _permit = permit;
198                    if let Err(e) = session(&child, &config, id, socket, peer).await {
199                        Log::new(Some(&child.status_tx))
200                            .warn(format!("Cap'n Proto RPC connection {id} ended: {e:#}"));
201                    }
202                    child
203                        .state
204                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
205                        .await;
206                    let _ = child.status_tx.send("__UPDATE_UI__".into());
207                })
208                .await;
209        }
210    });
211    state.register_server_task(server_id, accept).await;
212    Ok(addr)
213}
214
215fn outcome(ctx: &SpawnContext, id: ConnectionId, method: &str, decision: &str) {
216    let summary = format!("Cap'n Proto RPC connection {id} method={method} decision={decision}");
217    let log = Log::new(Some(&ctx.status_tx));
218    if decision.starts_with("fail_closed_") || decision == "model_silent" {
219        log.error(summary);
220    } else {
221        log.info(summary);
222    }
223}
224
225pub async fn send<W: AsyncWrite + Unpin>(writer: &mut W, words: &[u64]) -> Result<usize> {
226    let bytes = layout::frame(words);
227    tokio::time::timeout(WRITE_TIMEOUT, writer.write_all(&bytes))
228        .await
229        .context("write deadline")??;
230    Ok(bytes.len())
231}
232
233/// The handler's answer to one call, as a Return message.
234async fn answer_call(
235    ctx: &SpawnContext,
236    config: &Config,
237    id: ConnectionId,
238    peer: SocketAddr,
239    (question, interface, method): (u32, u64, u16),
240    content: Target<'_>,
241) -> Result<Vec<u64>> {
242    let methods = config.schema.methods_of(config.interface);
243    let Some((_, iface, m)) = methods
244        .iter()
245        .find(|(i, _, m)| *i == interface && m.id == method)
246    else {
247        return rpc::return_exception(
248            question,
249            &format!("method @{method} of interface {interface:#x} is not implemented"),
250            rpc::EXC_UNIMPLEMENTED,
251        );
252    };
253    let params = match content {
254        Target::Struct(s) => config.schema.to_json(m.params, s)?,
255        Target::Null => json!({}),
256        _ => bail!("call parameters are not a struct"),
257    };
258    let event = Event::new(
259        &actions::CALL_EVENT,
260        json!({"interface": iface.name, "method": m.name, "params": params,
261               "results_shape": config.schema.describe(m.results), "remote_addr": peer.to_string()}),
262    );
263    let failed = |e: Option<&anyhow::Error>| {
264        let text = match e {
265            Some(e) => crate::utils::wire_failure::prefixed_wire_failure_text(e),
266            None => crate::utils::WireFailure::Unavailable.prefixed_text(),
267        };
268        let kind = if e.is_some_and(crate::llm::is_overload_error) {
269            rpc::EXC_OVERLOADED
270        } else {
271            rpc::EXC_FAILED
272        };
273        rpc::return_exception(question, text, kind)
274    };
275    let result = match call_llm(
276        &ctx.llm_client,
277        &ctx.state,
278        ctx.server_id,
279        Some(id),
280        &event,
281        &actions::CapnpRpcProtocol,
282    )
283    .await
284    {
285        Ok(r) => r,
286        Err(e) => {
287            outcome(ctx, id, &m.name, "fail_closed_llm_error");
288            return failed(Some(&e));
289        }
290    };
291    if !result.failures.is_empty() {
292        outcome(ctx, id, &m.name, "fail_closed_invalid_reply");
293        return failed(None);
294    }
295    let mut answers = Vec::new();
296    let mut pending = result.protocol_results;
297    while let Some(r) = pending.pop() {
298        match r {
299            ActionResult::Custom { name, data } if name.starts_with("capnp_") => {
300                answers.push((name, data))
301            }
302            ActionResult::Multiple(items) => pending.extend(items),
303            _ => {}
304        }
305    }
306    match answers.as_slice() {
307        [(name, data)] if name == "capnp_return" => {
308            let (dw, pc) = config.schema.struct_size(m.results)?;
309            let built = rpc::return_results(question, |b, payload| {
310                let s = b.init_struct(payload, 0, dw, pc)?;
311                config.schema.from_json(b, m.results, s, &data["results"])
312            });
313            match built {
314                Ok(words) => {
315                    outcome(ctx, id, &m.name, "model_answer");
316                    Ok(words)
317                }
318                Err(e) => {
319                    Log::new(Some(&ctx.status_tx)).error(format!(
320                        "Cap'n Proto RPC results for {} do not fit the schema: {e:#}",
321                        m.name
322                    ));
323                    outcome(ctx, id, &m.name, "fail_closed_invalid_reply");
324                    failed(None)
325                }
326            }
327        }
328        [(name, data)] if name == "capnp_exception" => {
329            let (reason, kind) = actions::check_exception(data)?;
330            outcome(ctx, id, &m.name, "model_reject");
331            rpc::return_exception(question, &reason, kind)
332        }
333        [] => {
334            outcome(ctx, id, &m.name, "model_silent");
335            failed(None)
336        }
337        _ => {
338            outcome(ctx, id, &m.name, "fail_closed_invalid_reply");
339            failed(None)
340        }
341    }
342}
343
344async fn session(
345    ctx: &SpawnContext,
346    config: &Config,
347    id: ConnectionId,
348    socket: TcpStream,
349    peer: SocketAddr,
350) -> Result<()> {
351    let (mut reader, mut writer) = tokio::io::split(socket);
352    // Bootstrap questions whose answer (the exported capability) may still be pipelined on.
353    let mut bootstraps: HashSet<u32> = HashSet::new();
354    while let Some(segments) = layout::read_message(&mut reader, config.idle).await? {
355        ctx.state
356            .update_connection_stats(ctx.server_id, id, None, None, Some(1), None)
357            .await;
358        let msg = Message::new(segments);
359        let reply = match rpc::decode(&msg)? {
360            Incoming::Bootstrap { question } => {
361                ensure!(
362                    bootstraps.len() < MAX_OPEN_BOOTSTRAPS,
363                    "more than {MAX_OPEN_BOOTSTRAPS} open bootstrap questions"
364                );
365                bootstraps.insert(question);
366                Some(rpc::return_capability(question, 0)?)
367            }
368            Incoming::Call {
369                question,
370                target,
371                interface,
372                method,
373                content,
374            } => {
375                let on_bootstrap = match target {
376                    CallTarget::Imported(0) => true,
377                    CallTarget::Promised { question, steps: 0 } => bootstraps.contains(&question),
378                    _ => false,
379                };
380                Some(if on_bootstrap {
381                    answer_call(
382                        ctx,
383                        config,
384                        id,
385                        peer,
386                        (question, interface, method),
387                        content,
388                    )
389                    .await?
390                } else {
391                    rpc::return_exception(
392                        question,
393                        "no capability at that target: only the bootstrap capability is exported",
394                        rpc::EXC_FAILED,
395                    )?
396                })
397            }
398            Incoming::Finish { question } => {
399                bootstraps.remove(&question);
400                None
401            }
402            // NetGet asks nothing and holds no imports, so there is nothing to act on.
403            Incoming::Return { .. }
404            | Incoming::Release
405            | Incoming::Resolve
406            | Incoming::Unimplemented => None,
407            Incoming::Abort { reason } => {
408                Log::new(Some(&ctx.status_tx)).info(format!(
409                    "Cap'n Proto RPC connection {id} aborted by the peer: {}",
410                    crate::utils::truncate::truncate_for_log(&reason, 256)
411                ));
412                return Ok(());
413            }
414            Incoming::Other(_) => Some(rpc::unimplemented(&msg)?),
415        };
416        if let Some(words) = reply {
417            let n = send(&mut writer, &words).await?;
418            ctx.state
419                .update_connection_stats(ctx.server_id, id, None, Some(n as u64), None, Some(1))
420                .await;
421        }
422    }
423    Ok(())
424}
netget/server/capnp_rpc/mod.rs

netget/server/capnp_rpc/
mod.rs
