superfluid_agent/
client.rs

1//! The Link W client.
2
3use std::collections::{HashMap, VecDeque};
4use std::io::{Read, Write};
5use std::os::unix::net::UnixStream;
6use std::time::Duration;
7
8use superfluid_proto::envelope::FrameClass;
9use superfluid_proto::handshake::{Hello, HelloAckW, LinkRole};
10use superfluid_proto::linkw::{self, msg_type};
11use superfluid_proto::{encode_frame, Frame, FrameDecoder};
12use superfluid_shm::{
13    fdpass, ControlSegment, RingRole, SharedRing, SharedRings, ShmSegment, CONTROL_BYTES, CONTROL_TAG,
14    LOGITS_RING, TOKEN_RING_IN, TOKEN_RING_OUT,
15};
16
17use crate::AgentError;
18
19pub const PROTO_VERSION: u16 = 8;
20
21pub struct WorkerClient {
22    frames: UnixStream,
23    fd_channel: UnixStream,
24    decoder: FrameDecoder<fn(u16) -> bool>,
25    seq: u64,
26    pub hello: HelloAckW,
27    rings: SharedRings,
28    bufs: HashMap<u64, ShmSegment>,
29    next_buf_id: u64,
30    pub op_done: VecDeque<linkw::StateOpDoneMsg>,
31}
32
33impl WorkerClient {
34    pub fn connect(frames: UnixStream, fd_channel: UnixStream) -> Result<WorkerClient, AgentError> {
35        frames.set_read_timeout(Some(Duration::from_millis(50)))?;
36        let mut client = WorkerClient {
37            frames,
38            fd_channel,
39            decoder: FrameDecoder::new(linkw::is_known_type as fn(u16) -> bool),
40            seq: 1,
41            hello: HelloAckW {
42                chosen_version: 0,
43                engine_bundle_hash: [0; 32],
44                capabilities: String::new(),
45                state_space_descriptors: Vec::new(),
46                registered_strategies: Vec::new(),
47                limits: linkw::LinkWLimits {
48                    max_frame: 0,
49                    op_window: 0,
50                    ring_specs: Vec::new(),
51                },
52                kv_bits: 0,
53            },
54            rings: SharedRings {
55                token_in: SharedRing::create(TOKEN_RING_IN, 1, 32, RingRole::Writer)?,
56                token_out: SharedRing::create(TOKEN_RING_OUT, 1, 32, RingRole::Reader)?,
57                logits: SharedRing::create(LOGITS_RING, 1, 32, RingRole::Reader)?,
58                control: None,
59            },
60            bufs: HashMap::new(),
61            next_buf_id: 1,
62            op_done: VecDeque::new(),
63        };
64
65        let hello = Hello {
66            proto_versions: (PROTO_VERSION, PROTO_VERSION),
67            link_role: LinkRole::NodeAgent,
68            auth: Vec::new(),
69        };
70        let ack: HelloAckW = client.request(msg_type::HELLO, &hello, msg_type::HELLO_ACK_W)?;
71        client.hello = ack;
72
73        for spec in client.hello.limits.ring_specs.clone() {
74            let role = if spec.ring_id == TOKEN_RING_IN {
75                RingRole::Writer
76            } else {
77                RingRole::Reader
78            };
79            let ring = SharedRing::create(spec.ring_id, spec.slots, spec.slot_bytes, role)?;
80            let attach = linkw::RingAttachMsg {
81                ring_id: spec.ring_id,
82                layout: spec,
83            };
84            let payload = postcard::to_stdvec(&attach)?;
85            let sent_seq = client.send(msg_type::RING_ATTACH, 0, &payload)?;
86            fdpass::send_fd(
87                client.frames_fd_channel(),
88                ring.raw_fd(),
89                spec.ring_id as u64,
90            )?;
91            let ack_frame = client.wait_correlated(sent_seq)?;
92            let _ack: linkw::RingAckMsg = postcard::from_bytes(&ack_frame.payload)?;
93            match spec.ring_id {
94                TOKEN_RING_IN => client.rings.token_in = ring,
95                TOKEN_RING_OUT => client.rings.token_out = ring,
96                LOGITS_RING => client.rings.logits = ring,
97                _ => return Err(AgentError::Protocol("worker advertised unknown ring id")),
98            }
99        }
100
101        let control = ControlSegment::create()?;
102        let payload = postcard::to_stdvec(&linkw::ControlAttachMsg { bytes: CONTROL_BYTES as u32 })?;
103        let sent_seq = client.send(msg_type::CONTROL_ATTACH, 0, &payload)?;
104        fdpass::send_fd(client.frames_fd_channel(), control.raw_fd(), CONTROL_TAG)?;
105        let ack_frame = client.wait_correlated(sent_seq)?;
106        let _ack: linkw::ControlAckMsg = postcard::from_bytes(&ack_frame.payload)?;
107        client.rings.control = Some(control);
108        Ok(client)
109    }
110
111    /// The daemon's control words for this worker (end the running tick at
112    /// the next step boundary; prefill in short steps), settable from any
113    /// thread.
114    pub fn yield_signal(&self) -> Option<superfluid_shm::YieldSignal> {
115        self.rings.control.as_ref()?.signal().ok()
116    }
117
118    fn frames_fd_channel(&self) -> &UnixStream {
119        &self.fd_channel
120    }
121
122    pub fn raw_frames_stream(&mut self) -> &mut UnixStream {
123        &mut self.frames
124    }
125
126    fn send(&mut self, msg_type: u16, correlation: u64, payload: &[u8]) -> Result<u64, AgentError> {
127        let class = if msg_type & 0x8000 != 0 {
128            FrameClass::Droppable
129        } else {
130            FrameClass::Required
131        };
132        let seq = self.seq;
133        let bytes = encode_frame(PROTO_VERSION, msg_type, seq, correlation, class, payload)?;
134        self.seq += 1;
135        self.frames.write_all(&bytes)?;
136        Ok(seq)
137    }
138
139    fn wait_correlated(&mut self, sent_seq: u64) -> Result<Frame, AgentError> {
140        self.wait_correlated_streaming(sent_seq, &mut |_| true)
141    }
142
143    fn wait_correlated_emitting(
144        &mut self,
145        sent_seq: u64,
146        on_emit: &mut dyn FnMut(linkw::TickEmitMsg),
147    ) -> Result<Frame, AgentError> {
148        let mut buf = [0u8; 64 * 1024];
149        loop {
150            while let Some(frame) = self.decoder.decode_next()? {
151                if self.hello.chosen_version != 0
152                    && frame.proto_version != self.hello.chosen_version
153                {
154                    return Err(AgentError::Protocol(
155                        "frame version differs from the negotiated version",
156                    ));
157                }
158                if frame.correlation_id == sent_seq {
159                    return Ok(frame);
160                }
161                if frame.msg_type == msg_type::TICK_EMIT {
162                    match postcard::from_bytes(&frame.payload) {
163                        Ok(em) => on_emit(em),
164                        Err(e) => {
165                            tracing::warn!(error = %e, "dropping malformed TICK_EMIT frame")
166                        }
167                    }
168                    continue;
169                }
170                self.note_async(frame)?;
171            }
172            match self.frames.read(&mut buf) {
173                Ok(0) => return Err(AgentError::Closed),
174                Ok(n) => self.decoder.feed(&buf[..n]),
175                Err(e)
176                    if e.kind() == std::io::ErrorKind::WouldBlock
177                        || e.kind() == std::io::ErrorKind::TimedOut =>
178                {
179                    continue;
180                }
181                Err(e) => return Err(e.into()),
182            }
183        }
184    }
185
186    fn wait_correlated_streaming(
187        &mut self,
188        sent_seq: u64,
189        on_segment: &mut dyn FnMut(linkw::TranscribeSegmentMsg) -> bool,
190    ) -> Result<Frame, AgentError> {
191        let mut buf = [0u8; 64 * 1024];
192        let mut cancelled = false;
193        loop {
194            while let Some(frame) = self.decoder.decode_next()? {
195                if self.hello.chosen_version != 0
196                    && frame.proto_version != self.hello.chosen_version
197                {
198                    return Err(AgentError::Protocol(
199                        "frame version differs from the negotiated version",
200                    ));
201                }
202                if frame.correlation_id == sent_seq {
203                    return Ok(frame);
204                }
205                if frame.msg_type == msg_type::TRANSCRIBE_SEGMENT {
206                    if !on_segment(postcard::from_bytes(&frame.payload)?) && !cancelled {
207                        cancelled = true;
208                        let payload = postcard::to_stdvec(&linkw::TranscribeCancelMsg {})?;
209                        self.send(msg_type::TRANSCRIBE_CANCEL, 0, &payload)?;
210                    }
211                    continue;
212                }
213                self.note_async(frame)?;
214            }
215            match self.frames.read(&mut buf) {
216                Ok(0) => return Err(AgentError::Closed),
217                Ok(n) => self.decoder.feed(&buf[..n]),
218                Err(e)
219                    if e.kind() == std::io::ErrorKind::WouldBlock
220                        || e.kind() == std::io::ErrorKind::TimedOut => {}
221                Err(e) => return Err(e.into()),
222            }
223        }
224    }
225
226    fn note_async(&mut self, frame: Frame) -> Result<(), AgentError> {
227        if frame.msg_type == msg_type::STATE_OP_DONE {
228            let done: linkw::StateOpDoneMsg = postcard::from_bytes(&frame.payload)?;
229            self.op_done.push_back(done);
230        }
231        Ok(())
232    }
233
234    fn request<Req: serde::Serialize, Res: serde::de::DeserializeOwned>(
235        &mut self,
236        req_type: u16,
237        req: &Req,
238        res_type: u16,
239    ) -> Result<Res, AgentError> {
240        let payload = postcard::to_stdvec(req)?;
241        let sent = self.send(req_type, 0, &payload)?;
242        let frame = self.wait_correlated(sent)?;
243        if frame.msg_type != res_type {
244            return Err(AgentError::Protocol("unexpected response type"));
245        }
246        Ok(postcard::from_bytes(&frame.payload)?)
247    }
248
249    pub fn stage_prompt(&mut self, tokens: &[u32]) -> Result<linkw::TokenRefMsg, AgentError> {
250        let r = self.rings.stage_prompt(tokens)?;
251        Ok(linkw::TokenRefMsg {
252            ring_id: r.ring_id,
253            index: r.index,
254            count: r.count,
255            generation: r.generation,
256        })
257    }
258
259    pub fn read_tokens(&self, r: &linkw::TokenRefMsg) -> Result<Vec<u32>, AgentError> {
260        use superfluid_engine::Rings;
261        self.rings
262            .read_tokens(&superfluid_abi::TokenRef {
263                ring_id: r.ring_id,
264                index: r.index,
265                count: r.count,
266                _pad0: 0,
267                generation: r.generation,
268            })
269            .map_err(|_| AgentError::Protocol("stale emitted-token reference"))
270    }
271
272    pub fn read_logits(&self, r: &linkw::RingRefMsg) -> Result<Vec<f32>, AgentError> {
273        use superfluid_engine::Rings;
274        self.rings
275            .read_logits_row(&superfluid_abi::RingRef {
276                ring_id: r.ring_id,
277                index: r.index,
278                generation: r.generation,
279            })
280            .map_err(|_| AgentError::Protocol("stale logits reference"))
281    }
282
283    pub fn tick(&mut self, plan: linkw::TickPlanMsg) -> Result<linkw::TickEventsMsg, AgentError> {
284        let res: linkw::TickResultMsg = self.request(
285            msg_type::TICK_SUBMIT,
286            &linkw::TickSubmitMsg { plan },
287            msg_type::TICK_RESULT,
288        )?;
289        match res {
290            linkw::TickResultMsg::Events(ev) => Ok(ev),
291            linkw::TickResultMsg::Rejected { status } => Err(AgentError::Rejected(status)),
292        }
293    }
294
295    pub fn tick_streaming(
296        &mut self,
297        plan: linkw::TickPlanMsg,
298        on_emit: &mut dyn FnMut(linkw::TickEmitMsg),
299    ) -> Result<linkw::TickEventsMsg, AgentError> {
300        let payload = postcard::to_stdvec(&linkw::TickSubmitMsg { plan })?;
301        let sent = self.send(msg_type::TICK_SUBMIT, 0, &payload)?;
302        let frame = self.wait_correlated_emitting(sent, on_emit)?;
303        if frame.msg_type != msg_type::TICK_RESULT {
304            return Err(AgentError::Protocol("unexpected response type"));
305        }
306        match postcard::from_bytes::<linkw::TickResultMsg>(&frame.payload)? {
307            linkw::TickResultMsg::Events(ev) => Ok(ev),
308            linkw::TickResultMsg::Rejected { status } => Err(AgentError::Rejected(status)),
309        }
310    }
311
312    pub fn match_prefix(
313        &mut self,
314        space_id: u32,
315        span: linkw::TokenRefMsg,
316    ) -> Result<linkw::MatchResMsg, AgentError> {
317        self.request(
318            msg_type::MATCH_REQ,
319            &linkw::MatchReqMsg {
320                space_id,
321                span,
322                media_deps: Vec::new(),
323            },
324            msg_type::MATCH_RES,
325        )
326    }
327
328    pub fn seed_release(&mut self, seed_handle: u64) -> Result<(), AgentError> {
329        let payload = postcard::to_stdvec(&linkw::SeedReleaseMsg { seed_handle })?;
330        self.send(msg_type::SEED_RELEASE, 0, &payload)?;
331        Ok(())
332    }
333
334    pub fn seed_acquire(
335        &mut self,
336        span: linkw::TokenRefMsg,
337        prefix_len: u64,
338        determinism_class: u8,
339    ) -> Result<u64, AgentError> {
340        self.seed_acquire_leased(span, prefix_len, determinism_class)
341            .map(|(handle, _)| handle)
342    }
343
344    pub fn seed_acquire_leased(
345        &mut self,
346        span: linkw::TokenRefMsg,
347        prefix_len: u64,
348        determinism_class: u8,
349    ) -> Result<(u64, Option<u64>), AgentError> {
350        let payload = postcard::to_stdvec(&linkw::SeedAcquireMsg {
351            span,
352            prefix_len,
353            determinism_class,
354        })?;
355        let sent = self.send(msg_type::SEED_ACQUIRE, 0, &payload)?;
356        let frame = self.wait_correlated(sent)?;
357        if frame.msg_type != msg_type::SEED_GRANT {
358            return Err(AgentError::Protocol("unexpected response type"));
359        }
360        match linkw::SeedGrantMsg::decode(&frame.payload)? {
361            linkw::SeedGrantMsg::Granted {
362                seed_handle,
363                ttl_ticks,
364                ..
365            } => Ok((seed_handle, ttl_ticks)),
366            linkw::SeedGrantMsg::Refused { status } => Err(AgentError::Rejected(status)),
367        }
368    }
369
370    pub fn buf_create(&mut self, len: u64) -> Result<u64, AgentError> {
371        let buf_id = self.next_buf_id;
372        self.next_buf_id += 1;
373        let seg = ShmSegment::create(len as usize)?;
374        let alloc = postcard::to_stdvec(&linkw::BufAllocMsg { buf_id, len })?;
375        self.send(msg_type::BUF_ALLOC, 0, &alloc)?;
376        let attach = postcard::to_stdvec(&linkw::BufAttachMsg {
377            buf_id,
378            len,
379            generation: 1,
380        })?;
381        self.send(msg_type::BUF_ATTACH, 0, &attach)?;
382        fdpass::send_fd(&self.fd_channel, seg.raw_fd(), buf_id)?;
383        self.bufs.insert(buf_id, seg);
384        Ok(buf_id)
385    }
386
387    pub fn buf_read(&self, buf_id: u64, len: usize) -> Result<Vec<u8>, AgentError> {
388        let seg = self
389            .bufs
390            .get(&buf_id)
391            .ok_or(AgentError::Protocol("unknown buffer"))?;
392        let mut out = vec![0u8; len.min(seg.len())];
393        seg.read_at(0, &mut out)?;
394        Ok(out)
395    }
396
397    pub fn buf_write(&self, buf_id: u64, data: &[u8]) -> Result<(), AgentError> {
398        let seg = self
399            .bufs
400            .get(&buf_id)
401            .ok_or(AgentError::Protocol("unknown buffer"))?;
402        seg.write_at(0, data)?;
403        Ok(())
404    }
405
406    pub fn buf_release(&mut self, buf_id: u64) -> Result<(), AgentError> {
407        let payload = postcard::to_stdvec(&linkw::BufReleaseMsg { buf_id })?;
408        self.send(msg_type::BUF_RELEASE, 0, &payload)?;
409        self.bufs.remove(&buf_id);
410        Ok(())
411    }
412
413    pub fn state_op(&mut self, op: linkw::StateOpKind) -> Result<u64, AgentError> {
414        let res: linkw::StateOpAcceptedMsg = self.request(
415            msg_type::STATE_OP_SUBMIT,
416            &linkw::StateOpSubmitMsg { op },
417            msg_type::STATE_OP_ACCEPTED,
418        )?;
419        match res {
420            linkw::StateOpAcceptedMsg::Accepted { op_handle } => Ok(op_handle),
421            linkw::StateOpAcceptedMsg::Rejected { status } => Err(AgentError::Rejected(status)),
422        }
423    }
424
425    pub fn wait_op_done(&mut self, op: u64) -> Result<linkw::StateOpDoneMsg, AgentError> {
426        loop {
427            if let Some(pos) = self.op_done.iter().position(|d| d.op_handle == op) {
428                return Ok(self.op_done.remove(pos).expect("indexed"));
429            }
430            let mut buf = [0u8; 64 * 1024];
431            while let Some(frame) = self.decoder.decode_next()? {
432                self.note_async(frame)?;
433            }
434            match self.frames.read(&mut buf) {
435                Ok(0) => return Err(AgentError::Closed),
436                Ok(n) => self.decoder.feed(&buf[..n]),
437                Err(e)
438                    if e.kind() == std::io::ErrorKind::WouldBlock
439                        || e.kind() == std::io::ErrorKind::TimedOut => {}
440                Err(e) => return Err(e.into()),
441            }
442        }
443    }
444
445    pub fn state_sync(
446        &mut self,
447        req: linkw::StateSyncReqMsg,
448    ) -> Result<linkw::StateSyncOkMsg, AgentError> {
449        let res: linkw::StateSyncResMsg =
450            self.request(msg_type::STATE_SYNC_REQ, &req, msg_type::STATE_SYNC_RES)?;
451        match res {
452            linkw::StateSyncResMsg::Ok(ok) => Ok(ok),
453            linkw::StateSyncResMsg::Err { status } => Err(AgentError::Rejected(status)),
454        }
455    }
456
457    pub fn send_transcribe_cancel(&mut self) -> Result<(), AgentError> {
458        let payload = postcard::to_stdvec(&linkw::TranscribeCancelMsg {})?;
459        self.send(msg_type::TRANSCRIBE_CANCEL, 0, &payload)?;
460        Ok(())
461    }
462
463    pub fn state_sync_streaming(
464        &mut self,
465        req: linkw::StateSyncReqMsg,
466        on_segment: &mut dyn FnMut(linkw::TranscribeSegmentMsg) -> bool,
467    ) -> Result<linkw::StateSyncOkMsg, AgentError> {
468        let payload = postcard::to_stdvec(&req)?;
469        let sent = self.send(msg_type::STATE_SYNC_REQ, 0, &payload)?;
470        let frame = self.wait_correlated_streaming(sent, on_segment)?;
471        if frame.msg_type != msg_type::STATE_SYNC_RES {
472            return Err(AgentError::Protocol("unexpected response type"));
473        }
474        match postcard::from_bytes::<linkw::StateSyncResMsg>(&frame.payload)? {
475            linkw::StateSyncResMsg::Ok(ok) => Ok(ok),
476            linkw::StateSyncResMsg::Err { status } => Err(AgentError::Rejected(status)),
477        }
478    }
479
480    pub fn register_strategy(
481        &mut self,
482        reg: linkw::StrategyRegisterMsg,
483    ) -> Result<linkw::StrategyGrantOkMsg, AgentError> {
484        let res: linkw::StrategyGrantMsg =
485            self.request(msg_type::STRATEGY_REGISTER, &reg, msg_type::STRATEGY_GRANT)?;
486        match res {
487            linkw::StrategyGrantMsg::Granted(ok) => Ok(ok),
488            linkw::StrategyGrantMsg::Refused { status } => Err(AgentError::Rejected(status)),
489        }
490    }
491
492    pub fn ping(&mut self) -> Result<linkw::PongMsg, AgentError> {
493        let ping = linkw::PingMsg {
494            queue_depth: 0,
495            mem: linkw::MemCountersMsg {
496                allocated_bytes: 0,
497                host_retained_bytes: 0,
498                pool_blocks_total: 0,
499                pool_blocks_used: 0,
500                pool_bytes_evictable: 0,
501                tentative_bytes: 0,
502            },
503        };
504        self.request(msg_type::PING, &ping, msg_type::PONG)
505    }
506
507    pub fn drain(&mut self) -> Result<(), AgentError> {
508        let _: linkw::DrainedMsg =
509            self.request(msg_type::DRAIN_REQ, &linkw::DrainReqMsg, msg_type::DRAINED)?;
510        Ok(())
511    }
512}
superfluid_agent/client.rs

superfluid_agent/
client.rs
