regy 0.1.4

Private-by-default desktop agent for the Regy web interface
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

use crate::domain::{
    errors::{AgentError, AgentResult, ErrorCode},
    pi_rpc::PiRpcOutput,
};

pub(crate) struct PiRpcDecoder {
    buffer: Vec<u8>,
    max_record_bytes: usize,
}

pub(crate) struct SizedPiRpcOutput {
    pub(crate) output: PiRpcOutput,
    // Exact JSON bytes received, excluding the LF delimiter and the optional
    // CR stripped from CRLF framing.
    pub(crate) json_bytes: usize,
}

impl PiRpcDecoder {
    pub(crate) fn new(max_record_bytes: usize) -> Self {
        Self {
            buffer: Vec::new(),
            max_record_bytes,
        }
    }

    #[cfg(test)]
    pub(crate) fn push(&mut self, chunk: &[u8]) -> AgentResult<Vec<PiRpcOutput>> {
        Ok(self
            .push_sized(chunk)?
            .into_iter()
            .map(|record| record.output)
            .collect())
    }

    pub(crate) fn push_sized(&mut self, mut chunk: &[u8]) -> AgentResult<Vec<SizedPiRpcOutput>> {
        let mut records = Vec::new();

        while let Some(newline) = chunk.iter().position(|byte| *byte == b'\n') {
            self.append(&chunk[..newline])?;
            records.push(self.decode_buffered_record()?);
            chunk = &chunk[newline + 1..];
        }

        self.append(chunk)?;
        Ok(records)
    }

    pub(crate) fn finish(&mut self) -> AgentResult<Vec<PiRpcOutput>> {
        if self.buffer.iter().all(u8::is_ascii_whitespace) {
            self.buffer.clear();
            return Ok(Vec::new());
        }

        Err(invalid_message(
            "Pi RPC stream ended with an incomplete JSONL record",
        ))
    }

    fn append(&mut self, bytes: &[u8]) -> AgentResult<()> {
        let new_len = self.buffer.len().saturating_add(bytes.len());
        // The limit applies to JSON bytes. We temporarily allow exactly one extra
        // trailing CR so a maximum-sized record can still arrive with CRLF framing.
        let is_optional_cr = new_len == self.max_record_bytes.saturating_add(1)
            && bytes.last().or_else(|| self.buffer.last()) == Some(&b'\r');
        if new_len > self.max_record_bytes && !is_optional_cr {
            return Err(invalid_message(format!(
                "Pi RPC record exceeds {} byte limit",
                self.max_record_bytes
            )));
        }

        self.buffer.extend_from_slice(bytes);
        Ok(())
    }

    fn decode_buffered_record(&mut self) -> AgentResult<SizedPiRpcOutput> {
        if self.buffer.last() == Some(&b'\r') {
            self.buffer.pop();
        }
        if self.buffer.is_empty() {
            return Err(invalid_message("Pi RPC stream contained an empty record"));
        }

        let record = std::str::from_utf8(&self.buffer).map_err(|error| {
            invalid_message(format!("Pi RPC record is not valid UTF-8: {error}"))
        })?;
        let output = serde_json::from_str(record)
            .map_err(|error| invalid_message(format!("invalid Pi RPC JSON record: {error}")))?;
        let json_bytes = self.buffer.len();
        self.buffer.clear();
        Ok(SizedPiRpcOutput { output, json_bytes })
    }
}

fn invalid_message(message: impl Into<String>) -> AgentError {
    AgentError::new(ErrorCode::InvalidMessage, message)
}