mbx_cache_core/
client.rs

1//! Blocking client for the local cache-agent protocol.
2
3use crate::{AGENT_PROTOCOL_VERSION, AgentRequest, AgentResponse};
4use eyre::{Result, bail};
5use std::io::{BufRead, BufReader, Read, Write};
6
7/// A synchronous client over an already-connected local agent stream.
8///
9/// Embedders own endpoint discovery and platform transport creation; this type
10/// owns the version handshake and newline-delimited protocol framing.
11pub struct BlockingAgentClient<S> {
12    stream: BufReader<S>,
13}
14
15impl<S> BlockingAgentClient<S>
16where
17    S: Read + Write,
18{
19    /// Negotiate the exact agent protocol and application version.
20    pub fn connect(stream: S, client_version: impl Into<String>) -> Result<Self> {
21        let mut client = Self {
22            stream: BufReader::new(stream),
23        };
24        match client.request(AgentRequest::Hello {
25            protocol: AGENT_PROTOCOL_VERSION,
26            client_version: client_version.into(),
27        })? {
28            AgentResponse::Hello { protocol, .. } if protocol == AGENT_PROTOCOL_VERSION => {
29                Ok(client)
30            }
31            AgentResponse::Error { message } => bail!(message),
32            _ => bail!("cache agent returned an incompatible handshake"),
33        }
34    }
35
36    /// Send one request and wait for its response.
37    pub fn request(&mut self, request: AgentRequest) -> Result<AgentResponse> {
38        serde_json::to_writer(self.stream.get_mut(), &request)?;
39        self.stream.get_mut().write_all(b"\n")?;
40        self.stream.get_mut().flush()?;
41        let mut response = String::new();
42        if self.stream.read_line(&mut response)? == 0 {
43            bail!("cache agent closed the connection without a response");
44        }
45        Ok(serde_json::from_str(&response)?)
46    }
47
48    /// Begin a prediction-manifest run and return its opaque identifier.
49    pub fn begin_task(&mut self, task: impl Into<String>) -> Result<String> {
50        match self.request(AgentRequest::BeginTask { task: task.into() })? {
51            AgentResponse::TaskBegun { run } => Ok(run),
52            AgentResponse::Error { message } => bail!(message),
53            _ => bail!("cache agent returned an unexpected begin-task response"),
54        }
55    }
56
57    /// Commit predictions collected for a task run.
58    pub fn commit_task(&mut self, run: impl Into<String>) -> Result<()> {
59        match self.request(AgentRequest::CommitTask { run: run.into() })? {
60            AgentResponse::TaskCommitted => Ok(()),
61            AgentResponse::Error { message } => bail!(message),
62            _ => bail!("cache agent returned an unexpected commit-task response"),
63        }
64    }
65
66    /// Recover the connected stream.
67    pub fn into_inner(self) -> S {
68        self.stream.into_inner()
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use std::io::{Cursor, Result as IoResult};
76
77    struct ScriptedStream {
78        responses: Cursor<Vec<u8>>,
79        requests: Vec<u8>,
80    }
81
82    impl Read for ScriptedStream {
83        fn read(&mut self, buffer: &mut [u8]) -> IoResult<usize> {
84            self.responses.read(buffer)
85        }
86    }
87
88    impl Write for ScriptedStream {
89        fn write(&mut self, buffer: &[u8]) -> IoResult<usize> {
90            self.requests.extend_from_slice(buffer);
91            Ok(buffer.len())
92        }
93
94        fn flush(&mut self) -> IoResult<()> {
95            Ok(())
96        }
97    }
98
99    fn stream(responses: &[AgentResponse]) -> ScriptedStream {
100        let mut bytes = Vec::new();
101        for response in responses {
102            serde_json::to_writer(&mut bytes, response).unwrap();
103            bytes.push(b'\n');
104        }
105        ScriptedStream {
106            responses: Cursor::new(bytes),
107            requests: Vec::new(),
108        }
109    }
110
111    #[test]
112    fn begins_and_commits_task_runs() {
113        let responses = [
114            AgentResponse::Hello {
115                protocol: AGENT_PROTOCOL_VERSION,
116                agent_version: "0.5.1".into(),
117            },
118            AgentResponse::TaskBegun {
119                run: "run-1".into(),
120            },
121            AgentResponse::TaskCommitted,
122        ];
123        let mut client = BlockingAgentClient::connect(stream(&responses), "0.5.1").unwrap();
124        let run = client.begin_task("task-1").unwrap();
125        client.commit_task(&run).unwrap();
126        assert_eq!(run, "run-1");
127
128        let written = String::from_utf8(client.into_inner().requests).unwrap();
129        assert!(written.contains("\"type\":\"begin_task\""));
130        assert!(written.contains("\"type\":\"commit_task\""));
131    }
132
133    #[test]
134    fn rejects_a_different_protocol_version() {
135        let responses = [AgentResponse::Hello {
136            protocol: AGENT_PROTOCOL_VERSION + 1,
137            agent_version: "0.5.1".into(),
138        }];
139        let error = BlockingAgentClient::connect(stream(&responses), "0.5.1")
140            .err()
141            .expect("protocol mismatch should fail");
142        assert!(error.to_string().contains("incompatible handshake"));
143    }
144}
mbx_cache_core/client.rs

mbx_cache_core/
client.rs
