ts_wire/
frames.rs

1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! ZMQ message framing on the frontend–engine sockets.
22//!
23//! * Input (frontend ROUTER to engine DEALER): `[type, payload, aux...]`.
24//!   `type` is one byte so the engine dispatches without decoding: ADD
25//!   carries a [`TokenizedGenerateReqInput`] whose tensors may reference the
26//!   `aux` frames; ABORT carries a msgpack list of request ids; CONTROL
27//!   carries a [`ControlRequest`].
28//! * Output (engine PUSH to frontend PULL): a [`BatchTokenIDOutSlim`]
29//!   message, a [`ControlOutput`] (told apart by its tag), or the single
30//!   frame [`ENGINE_CORE_DEAD`] when the engine exits.
31
32use bytes::Bytes;
33
34use crate::{
35    codec::{decode_frames, encode_frames},
36    control::{ControlOutput, ControlRequest},
37    error::{Error, Result},
38    msgpack::Decoder,
39    output::BatchTokenIDOutSlim,
40    request::TokenizedGenerateReqInput,
41};
42
43/// Type byte of an ADD message.
44pub const REQ_TYPE_ADD: u8 = 0x00;
45/// Type byte of an ABORT message.
46pub const REQ_TYPE_ABORT: u8 = 0x01;
47/// Type byte of a CONTROL message.
48pub const REQ_TYPE_CONTROL: u8 = 0x02;
49/// The frame an engine sends on its output socket when it shuts down; the
50/// frontend then treats the engine as dead.
51pub const ENGINE_CORE_DEAD: &[u8] = b"ENGINE_CORE_DEAD";
52
53/// The ZMQ routing identity of engine `engine_index`: its index as two
54/// little-endian bytes.
55pub fn engine_identity(engine_index: u16) -> [u8; 2] {
56    engine_index.to_le_bytes()
57}
58
59/// A message on the engine's input socket.
60#[derive(Debug, Clone, PartialEq)]
61pub enum EngineRequest {
62    /// Start a request.
63    Add(Box<TokenizedGenerateReqInput>),
64    /// Abort requests by id.
65    Abort(Vec<String>),
66    /// A control operation.
67    Control(Box<ControlRequest>),
68}
69
70impl EngineRequest {
71    /// The message's frames.
72    pub fn to_frames(&self) -> Vec<Bytes> {
73        let (type_frame, payload): (&'static [u8], _) = match self {
74            EngineRequest::Add(request) => (&[REQ_TYPE_ADD], encode_frames(request.as_ref())),
75            EngineRequest::Abort(rids) => (&[REQ_TYPE_ABORT], encode_frames(rids)),
76            EngineRequest::Control(request) => {
77                (&[REQ_TYPE_CONTROL], encode_frames(request.as_ref()))
78            }
79        };
80        let mut frames = Vec::with_capacity(1 + payload.len());
81        frames.push(Bytes::from_static(type_frame));
82        frames.extend(payload);
83        frames
84    }
85
86    /// Decode a message from its frames.
87    pub fn from_frames(frames: &[Bytes]) -> Result<Self> {
88        let [type_frame, payload @ ..] = frames else {
89            return Err(Error::Invalid("an input message has no frames".to_owned()));
90        };
91        if payload.is_empty() {
92            return Err(Error::Invalid(
93                "an input message needs a type frame and a payload".to_owned(),
94            ));
95        }
96        match type_frame.as_ref() {
97            [REQ_TYPE_ADD] => Ok(EngineRequest::Add(Box::new(decode_frames(payload)?))),
98            [REQ_TYPE_ABORT] => Ok(EngineRequest::Abort(decode_frames(payload)?)),
99            [REQ_TYPE_CONTROL] => Ok(EngineRequest::Control(Box::new(decode_frames(payload)?))),
100            other => Err(Error::Invalid(format!(
101                "unknown request type frame {other:02x?}"
102            ))),
103        }
104    }
105}
106
107/// A message on the engine's output socket.
108#[derive(Debug, Clone, PartialEq)]
109pub enum EngineOutput {
110    /// One step's outputs.
111    Batch(Box<BatchTokenIDOutSlim>),
112    /// A control operation's reply.
113    Control(Box<ControlOutput>),
114    /// The engine shut down.
115    Dead,
116}
117
118impl EngineOutput {
119    /// The message's frames.
120    pub fn to_frames(&self) -> Vec<Bytes> {
121        match self {
122            EngineOutput::Batch(batch) => encode_frames(batch.as_ref()),
123            EngineOutput::Control(output) => encode_frames(output.as_ref()),
124            EngineOutput::Dead => vec![Bytes::from_static(ENGINE_CORE_DEAD)],
125        }
126    }
127
128    /// Decode a message from its frames.
129    pub fn from_frames(frames: &[Bytes]) -> Result<Self> {
130        match frames {
131            [] => Err(Error::Invalid("an output message has no frames".to_owned())),
132            [only] if only.as_ref() == ENGINE_CORE_DEAD => Ok(EngineOutput::Dead),
133            _ if ControlOutput::knows(Decoder::new(frames).peek_tag()?) => {
134                Ok(EngineOutput::Control(Box::new(decode_frames(frames)?)))
135            }
136            _ => Ok(EngineOutput::Batch(Box::new(decode_frames(frames)?))),
137        }
138    }
139}
ts_wire/frames.rs

ts_wire/
frames.rs
