Skip to main content

netget/server/kafka/
mod.rs

1//! Kafka broker — Rust owns the wire format, the LLM owns the content.
2//!
3//! # What this broker supports
4//!
5//! A deliberately narrow subset of the Kafka protocol: enough for a client to complete
6//! its opening sequence (ApiVersions → Metadata) and then produce and fetch records.
7//!
8//! | API | key | versions | answered by |
9//! |---|---|---|---|
10//! | ApiVersions  | 18 | 0–3  | Rust (pure transport negotiation, no LLM call) |
11//! | Metadata     | 3  | 0–12 | LLM / script / static handler |
12//! | Produce      | 0  | 0–11 | LLM / script / static handler |
13//! | Fetch        | 1  | 0–16 | LLM / script / static handler |
14//! | OffsetCommit | 8  | 0–9  | LLM / script / static handler |
15//!
16//! Everything else — `ListOffsets`, `FindCoordinator`, `JoinGroup`, `SyncGroup`,
17//! `Heartbeat`, `OffsetFetch`, the admin APIs — is **not implemented**. Those keys are
18//! not advertised in ApiVersions, and if a client sends one anyway the connection is
19//! closed with an ERROR log rather than answered with a body this broker cannot build
20//! correctly. The practical consequence: a consumer must use manual partition
21//! assignment and an explicit fetch offset. Consumer *groups* do not work, because
22//! group coordination needs the APIs above.
23//!
24//! The ceilings are the newest versions `kafka-protocol` 0.14 knows, flexible (tagged-field)
25//! versions included; Apache Kafka's Java client and franz-go both complete at Metadata v12,
26//! Produce v11 and Fetch v16. From Fetch v13 a topic is named by id: ids are derived from the
27//! name ([`topic_id`]) because nothing is stored, and the broker remembers the ones it has
28//! reported so an id-only request can be resolved.
29//!
30//! # Version handling
31//!
32//! `api_key` and `api_version` are read from the first four bytes of the frame, the
33//! request header is decoded at `ApiKey::request_header_version(api_version)`, the body
34//! at `api_version`, and the response header at
35//! `ApiKey::response_header_version(api_version)`. Nothing is hardcoded to version 0.
36//! `correlation_id` sits at a fixed offset in every request header version, so the error
37//! paths can echo it without having decoded the header.
38//!
39//! # No storage
40//!
41//! This module keeps no topics, no log and no consumer offsets. Every produce, fetch,
42//! metadata and offset-commit answer comes from an action returned by the LLM, a script
43//! handler or a static handler. A Fetch returns exactly the records the model supplied
44//! for that request. There is no per-connection or per-server bookkeeping of offsets:
45//! the only state is `cluster_id`, `broker_id` and the advertised host, all fixed at
46//! startup.
47//!
48//! # Failure is refusal, never a default answer
49//!
50//! If the model (or handler) returns nothing usable, the client gets the correct
51//! response *type* carrying `UNKNOWN_SERVER_ERROR` (-1). A model that wants to refuse
52//! uses the `error_response` action and picks its own Kafka error code, which is a
53//! structurally different path from silence. Nothing is ever invented to keep a client
54//! happy — with one documented exception: the broker list in a Metadata response
55//! defaults to this server's own listen address, because that is transport truth the
56//! model has no way to know and no client can proceed without it.
57
58pub mod actions;
59
60/// Re-export of the wire library.
61///
62/// Integration tests decode this broker's responses with `kafka-protocol`'s own
63/// code-generated client-side decoders — the opposite direction from the encoders used
64/// here — which is not possible from `tests/` unless the crate is reachable through
65/// `netget`. It is not a dev-dependency and adding one would mean touching `Cargo.toml`.
66pub use kafka_protocol;
67
68use crate::llm::action_helper::call_llm;
69use crate::llm::actions::protocol_trait::ActionResult;
70use crate::llm::ollama_client::OllamaClient;
71use crate::logging::emit::Log;
72use crate::protocol::Event;
73use crate::server::connection::ConnectionId;
74use crate::server::KafkaProtocol;
75use crate::state::app_state::AppState;
76use actions::{
77    FETCH_REQUEST_EVENT, METADATA_REQUEST_EVENT, OFFSET_COMMIT_REQUEST_EVENT, PRODUCE_REQUEST_EVENT,
78};
79use anyhow::Result;
80use bytes::Bytes;
81use kafka_protocol::messages::{
82    ApiKey, ApiVersionsResponse, BrokerId, FetchRequest, FetchResponse, MetadataRequest,
83    MetadataResponse, OffsetCommitRequest, OffsetCommitResponse, ProduceRequest, ProduceResponse,
84    RequestHeader, ResponseHeader,
85};
86use kafka_protocol::protocol::{Decodable, Encodable, StrBytes};
87use kafka_protocol::records::{
88    Compression, Record, RecordBatchDecoder, RecordBatchEncoder, RecordEncodeOptions, TimestampType,
89};
90use serde_json::{json, Value};
91use std::net::SocketAddr;
92use std::sync::Arc;
93use tokio::io::{AsyncReadExt, AsyncWriteExt};
94use tokio::net::{TcpListener, TcpStream};
95use tokio::sync::{mpsc, Mutex};
96use tracing::{debug, error, warn};
97
98/// How long to wait for a peer's first request after it has connected.
99///
100/// Kafka is client-speaks-first with no broker greeting, and every real client opens with
101/// `ApiVersions` before it does anything else — `librdkafka`, the Java client and `kcat` all
102/// send it inside the connect path. A peer that has connected and sent nothing has made no
103/// request at all, so it gets the short bound.
104const FIRST_REQUEST_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
105
106/// How long to wait for a *further* request once one has been answered.
107///
108/// Ten minutes, which is not a guess: it is `connections.max.idle.ms`, the broker setting real
109/// Kafka ships with exactly this job, at exactly this default. A producer between batches or a
110/// consumer between long polls is legitimately silent for minutes, and copying the broker's own
111/// number is the one choice no real client can be surprised by.
112const IDLE_BETWEEN_REQUESTS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600);
113
114/// Concurrent connections this server admits.
115///
116/// Each connection may buffer a request of up to [`MAX_REQUEST_BYTES`], so the cap is what
117/// turns that per-connection bound into a total one.
118const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
119
120/// A peer over [`MAX_CONNECTIONS`] gets a plain close, which is what a real broker does.
121///
122/// Every Kafka response is `(size)(correlation_id)(body)` and the correlation id belongs to the
123/// *request* — this peer has sent none, so no well-formed response exists to write. Real
124/// brokers hitting `max.connections` close the socket without a word for the same reason, and
125/// `librdkafka` reports it as a broker transport failure and backs off. The reason is in the
126/// log, tagged `decision=fail_closed_connection_cap`.
127const CONNECTION_CAP_REFUSAL: &[u8] = b"";
128
129/// Smallest useful Kafka request: api_key (i16) + api_version (i16) + correlation_id (i32).
130const MIN_REQUEST_BYTES: i32 = 8;
131
132/// Largest request accepted from a client. Real brokers cap this with
133/// `socket.request.max.bytes` (default 100 MiB); without a cap the wire-supplied
134/// size is an allocation primitive for any unauthenticated peer.
135pub const MAX_REQUEST_BYTES: usize = 100 * 1024 * 1024;
136
137/// Cap on how much of a request is hex-dumped at TRACE level.
138const MAX_TRACE_HEX_BYTES: usize = 4096;
139
140/// Upper bound on a partition index this broker will talk about. Partition indices
141/// arrive as unvalidated i32s and are used to size and index collections.
142const MAX_PARTITIONS: i32 = 1024;
143
144/// Upper bound on how many (topic, partition) units one request may contain. Each unit
145/// costs one LLM round trip, so an attacker naming 100k partitions in one Produce would
146/// otherwise be a request amplifier against the model backend.
147const MAX_UNITS_PER_REQUEST: usize = 64;
148
149/// Upper bound on records encoded into one Fetch response partition.
150const MAX_RECORDS_PER_FETCH: usize = 1000;
151
152/// Upper bound on the total record payload encoded into one Fetch response partition.
153const MAX_FETCH_PAYLOAD_BYTES: usize = 8 * 1024 * 1024;
154
155/// How many produced records are described to the model in one `kafka_produce_request`
156/// event, and how many bytes of each. Prompts, not wire limits.
157const MAX_EVENT_RECORDS: usize = 20;
158const MAX_EVENT_VALUE_BYTES: usize = 1024;
159
160// Kafka error codes used by this module. See
161// <https://kafka.apache.org/protocol.html#protocol_error_codes>.
162const ERR_UNKNOWN_SERVER_ERROR: i16 = -1;
163const ERR_NONE: i16 = 0;
164const ERR_CORRUPT_MESSAGE: i16 = 2;
165const ERR_UNKNOWN_TOPIC_OR_PARTITION: i16 = 3;
166/// Retriable in every Kafka client. Used when the *backend* is saturated rather than
167/// when anything about the request was wrong, so the client backs off and tries again
168/// instead of recording a permanent fault - the distinction `WireFailure` exists to draw.
169const ERR_REQUEST_TIMED_OUT: i16 = 7;
170const ERR_UNSUPPORTED_VERSION: i16 = 35;
171/// A topic named by an id this broker has never reported (Fetch v13+, Metadata v12).
172const ERR_UNKNOWN_TOPIC_ID: i16 = 100;
173
174/// Topic ids remembered for resolving Fetch v13+ and Metadata v12 requests, which name
175/// topics by id alone. Every id this broker hands out is derived from the name, so a
176/// forgotten one costs a client one `UNKNOWN_TOPIC_ID` and a metadata refresh, never a wrong
177/// answer. Past this the oldest one is forgotten, one at a time.
178pub const MAX_KNOWN_TOPIC_IDS: usize = 4096;
179
180/// Kafka's own limit on a topic name, in bytes (`Topic.MAX_NAME_LENGTH`). A Metadata answer
181/// naming a longer topic leaves that topic out, and its name is never remembered.
182pub const MAX_TOPIC_NAME_LEN: usize = 249;
183
184/// The topic ids this broker has reported, oldest first, so a request naming a topic by id
185/// alone can be answered. Holds at most [`MAX_KNOWN_TOPIC_IDS`] names of at most
186/// [`MAX_TOPIC_NAME_LEN`] bytes each.
187#[derive(Default)]
188pub struct TopicIds {
189    inner: std::sync::Mutex<TopicIdTable>,
190}
191
192#[derive(Default)]
193struct TopicIdTable {
194    names: std::collections::HashMap<uuid::Uuid, String>,
195    order: std::collections::VecDeque<uuid::Uuid>,
196}
197
198impl TopicIds {
199    /// Remember and return the id reported for `name`; `None`, remembering nothing, for a
200    /// name over [`MAX_TOPIC_NAME_LEN`] bytes. At [`MAX_KNOWN_TOPIC_IDS`] the oldest id is
201    /// forgotten to make room.
202    pub fn report(&self, name: &str) -> Option<uuid::Uuid> {
203        if name.len() > MAX_TOPIC_NAME_LEN {
204            return None;
205        }
206        let id = topic_id(name);
207        let mut table = self.inner.lock().unwrap_or_else(|p| p.into_inner());
208        if !table.names.contains_key(&id) {
209            while table.names.len() >= MAX_KNOWN_TOPIC_IDS {
210                let Some(oldest) = table.order.pop_front() else {
211                    break;
212                };
213                table.names.remove(&oldest);
214            }
215            table.order.push_back(id);
216        }
217        table.names.insert(id, name.to_string());
218        Some(id)
219    }
220
221    /// The topic a reported id stands for.
222    pub fn name(&self, id: &uuid::Uuid) -> Option<String> {
223        self.inner
224            .lock()
225            .unwrap_or_else(|p| p.into_inner())
226            .names
227            .get(id)
228            .cloned()
229    }
230
231    /// Ids remembered.
232    pub fn len(&self) -> usize {
233        self.inner
234            .lock()
235            .unwrap_or_else(|p| p.into_inner())
236            .names
237            .len()
238    }
239
240    pub fn is_empty(&self) -> bool {
241        self.len() == 0
242    }
243}
244
245/// The id this broker reports for a topic (Metadata v10+, Fetch v13+).
246///
247/// NetGet keeps no topics, so there is nowhere to store an id assigned at creation; it is
248/// derived from the name instead — two 64-bit FNV-1a hashes with different offsets, the
249/// RFC 9562 version nibble set to 8 (custom). The same name always gets the same id, across
250/// connections and restarts, which is what a real broker's stored id gives a client. The
251/// nil id and Kafka's reserved `0…01` are never produced.
252pub fn topic_id(name: &str) -> uuid::Uuid {
253    let fnv = |offset: u64| {
254        name.bytes().fold(offset, |h, b| {
255            (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3)
256        })
257    };
258    let mut bytes = [0u8; 16];
259    bytes[..8].copy_from_slice(&fnv(0xcbf2_9ce4_8422_2325).to_be_bytes());
260    bytes[8..].copy_from_slice(&fnv(0x6c62_272e_07bb_0142).to_be_bytes());
261    bytes[6] = (bytes[6] & 0x0f) | 0x80;
262    bytes[8] = (bytes[8] & 0x3f) | 0x80;
263    uuid::Uuid::from_bytes(bytes)
264}
265
266/// The API keys this broker implements, with the version range it implements for each.
267///
268/// This table is the single source of truth: it is what ApiVersions advertises and what
269/// the dispatcher validates against, so the two can never disagree.
270pub const SUPPORTED_APIS: &[(i16, i16, i16)] = &[
271    (ApiKey::Produce as i16, 0, 11),
272    (ApiKey::Fetch as i16, 0, 16),
273    (ApiKey::Metadata as i16, 0, 12),
274    (ApiKey::OffsetCommit as i16, 0, 9),
275    (ApiKey::ApiVersions as i16, 0, 3),
276];
277
278/// Version range this broker implements for `api_key`, if any.
279pub fn supported_versions(api_key: i16) -> Option<(i16, i16)> {
280    SUPPORTED_APIS
281        .iter()
282        .find(|(k, _, _)| *k == api_key)
283        .map(|(_, min, max)| (*min, *max))
284}
285
286/// Kafka broker server state.
287///
288/// Deliberately tiny: everything here is fixed at startup and describes *this process*,
289/// not the contents of any topic. There is no log and no offset store; see the module
290/// docs.
291pub struct KafkaServer {
292    /// Cluster ID reported in Metadata (v2+ only; earlier versions have no such field).
293    cluster_id: String,
294    /// Broker ID reported in Metadata and used as the default partition leader.
295    broker_id: i32,
296    /// Host advertised to clients in Metadata when the model does not name one.
297    advertised_host: String,
298    /// Ids this broker has reported, so a request naming a topic by id alone can be
299    /// answered.
300    topic_ids: TopicIds,
301}
302
303impl KafkaServer {
304    /// Remember (and return) the id reported for `name`; `None` for a name Kafka would refuse.
305    fn report_topic(&self, name: &str) -> Option<uuid::Uuid> {
306        self.topic_ids.report(name)
307    }
308
309    /// The topic a reported id stands for.
310    fn topic_name(&self, id: &uuid::Uuid) -> Option<String> {
311        self.topic_ids.name(id)
312    }
313}
314
315/// What the model said about one request unit.
316enum Reply {
317    /// The expected response action, with its parameters.
318    Data(Value),
319    /// An explicit `error_response` with a model-chosen Kafka error code.
320    Error(i16),
321    /// Nothing usable came back. Never treated as approval.
322    ///
323    /// Carries the error code to report, because the two ways of getting here are not the
324    /// same thing to a client: a saturated backend is transient and gets the retriable
325    /// `REQUEST_TIMED_OUT`, while a handler that ran and produced nothing usable gets
326    /// `UNKNOWN_SERVER_ERROR`, which clients treat as permanent. Collapsing both onto -1
327    /// made an overloaded backend look like a broken broker.
328    NoAnswer(i16),
329}
330
331/// The write half of one live connection.
332///
333/// Shared between the session below and the dashboard's peer-command task
334/// (`server::peer_support`), so an injected send and the broker's own reply cannot interleave
335/// bytes into the middle of each other's frame. Split from the socket with
336/// `tokio::io::split`; never cloned.
337type SharedWrite = Arc<Mutex<tokio::io::WriteHalf<TcpStream>>>;
338
339impl KafkaServer {
340    /// Spawn Kafka broker with LLM integration
341    pub async fn spawn_with_llm_actions(
342        listen_addr: SocketAddr,
343        llm_client: OllamaClient,
344        app_state: Arc<AppState>,
345        status_tx: mpsc::UnboundedSender<String>,
346        server_id: crate::state::ServerId,
347        startup_params: Option<crate::protocol::StartupParams>,
348    ) -> Result<SocketAddr> {
349        let cluster_id = startup_params
350            .as_ref()
351            .map(|p| p.get_optional_string("cluster_id"))
352            .transpose()?
353            .flatten()
354            .unwrap_or_else(|| "netget-kafka-1".to_string());
355        let broker_id = startup_params
356            .as_ref()
357            .map(|p| p.get_optional_i64("broker_id"))
358            .transpose()?
359            .flatten()
360            .unwrap_or(0)
361            .clamp(0, i32::MAX as i64) as i32;
362        let advertised_host_param = startup_params
363            .as_ref()
364            .map(|p| p.get_optional_string("advertised_host"))
365            .transpose()?
366            .flatten();
367
368        let listener = TcpListener::bind(listen_addr).await?;
369        let local_addr = listener.local_addr()?;
370
371        // A client connects to whatever host the Metadata response names. If we are
372        // bound to a wildcard address, the bound IP is not connectable, so fall back to
373        // localhost rather than advertising 0.0.0.0.
374        let advertised_host = advertised_host_param.unwrap_or_else(|| {
375            if local_addr.ip().is_unspecified() {
376                "localhost".to_string()
377            } else {
378                local_addr.ip().to_string()
379            }
380        });
381
382        let log = Log::new(Some(&status_tx));
383        log.info(format!(
384            "Kafka broker listening on {} (cluster={}, broker_id={}, advertised_host={})",
385            local_addr, cluster_id, broker_id, advertised_host
386        ));
387        // Generated from SUPPORTED_APIS, so the log cannot disagree with what is advertised.
388        let supported = SUPPORTED_APIS
389            .iter()
390            .map(|(key, min, max)| match ApiKey::try_from(*key) {
391                Ok(api) => format!("{api:?} v{min}-{max}"),
392                Err(_) => format!("key {key} v{min}-{max}"),
393            })
394            .collect::<Vec<_>>()
395            .join(", ");
396        log.info(format!(
397            "Kafka supports {}. Other API keys (ListOffsets, FindCoordinator, JoinGroup, \
398             SyncGroup, Heartbeat, OffsetFetch, admin) are not implemented, so consumer groups \
399             do not work; consumers must assign partitions and fetch from an explicit offset. \
400             Broker advertises host '{}'.",
401            supported, advertised_host
402        ));
403
404        let server = Arc::new(KafkaServer {
405            cluster_id,
406            broker_id,
407            advertised_host,
408            topic_ids: Default::default(),
409        });
410
411        let protocol = Arc::new(KafkaProtocol::new());
412
413        let task_registrar = app_state.clone();
414        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
415        let accept_handle = tokio::spawn(async move {
416            loop {
417                match crate::server::accept_bounded::accept_bounded(
418                    &listener,
419                    &limiter,
420                    CONNECTION_CAP_REFUSAL,
421                    "Kafka",
422                    Some(&status_tx),
423                )
424                .await
425                {
426                    Ok((stream, peer_addr, permit)) => {
427                        Log::new(Some(&status_tx))
428                            .debug(format!("Kafka client connected from {}", peer_addr));
429
430                        let connection_id =
431                            ConnectionId::new(app_state.get_next_unified_id().await);
432                        let llm_clone = llm_client.clone();
433                        let state_clone = app_state.clone();
434                        let status_clone = status_tx.clone();
435                        let server_clone = server.clone();
436                        let protocol_clone = protocol.clone();
437
438                        use crate::state::server::{
439                            ConnectionState as ServerConnectionState, ConnectionStatus,
440                            ProtocolConnectionInfo,
441                        };
442                        let now = crate::utils::clock::Instant::now();
443                        let conn_state = ServerConnectionState {
444                            id: connection_id,
445                            remote_addr: peer_addr,
446                            local_addr,
447                            bytes_sent: 0,
448                            bytes_received: 0,
449                            packets_sent: 0,
450                            packets_received: 0,
451                            last_activity: now,
452                            status: ConnectionStatus::Active,
453                            status_changed_at: now,
454                            protocol_info: ProtocolConnectionInfo::empty(),
455                        };
456                        app_state
457                            .add_connection_to_server(server_id, conn_state)
458                            .await;
459                        let _ = status_clone.send("__UPDATE_UI__".to_string());
460
461                        // Tracked, not detached: stop_server must abort this task too.
462                        let task_owner = app_state.clone();
463                        task_owner
464                            .spawn_server_task(server_id, async move {
465                                // Held for the life of the connection, so the cap counts live
466                                // clients rather than accepts.
467                                let _permit = permit;
468                                let result = Self::handle_connection(
469                                    stream,
470                                    peer_addr,
471                                    local_addr,
472                                    connection_id,
473                                    server_clone,
474                                    llm_clone,
475                                    state_clone.clone(),
476                                    status_clone.clone(),
477                                    server_id,
478                                    protocol_clone,
479                                )
480                                .await;
481
482                                if let Err(e) = result {
483                                    error!("Kafka connection error: {}", e);
484                                }
485
486                                // Connections used to be added and never removed, so the
487                                // TUI accumulated Active entries for dead sockets.
488                                //
489                                // `close_connection_on_server`, not `update_connection_status`:
490                                // both mark the row `Closed`, but only the former also runs
491                                // `cleanup_connection_tasks`. A connection-scoped scheduled
492                                // task belongs to a connection that no longer exists, and a
493                                // recurring one keeps ticking after the socket is gone — each
494                                // tick an LLM prompt about a peer that has hung up. Marking the
495                                // row closed and leaving the tasks running is the defect
496                                // `remove_server` was restructured to prevent, one level down.
497                                state_clone
498                                    .close_connection_on_server(server_id, connection_id)
499                                    .await;
500                                let _ = status_clone.send("__UPDATE_UI__".to_string());
501                            })
502                            .await;
503                    }
504                    Err(e) => {
505                        // Without a break this spun a hot loop on a persistent
506                        // error (EMFILE), saturating a core and flooding the
507                        // unbounded status channel.
508                        Log::new(Some(&status_tx))
509                            .error(format!("Kafka accept error, loop stopping: {}", e));
510                        break;
511                    }
512                }
513            }
514        });
515
516        task_registrar
517            .register_server_task(server_id, accept_handle)
518            .await;
519
520        Ok(local_addr)
521    }
522
523    /// Set one connection up, run it, and tear it down on every exit path.
524    #[allow(clippy::too_many_arguments)]
525    async fn handle_connection(
526        stream: TcpStream,
527        peer_addr: SocketAddr,
528        local_addr: SocketAddr,
529        connection_id: ConnectionId,
530        server: Arc<KafkaServer>,
531        llm_client: OllamaClient,
532        app_state: Arc<AppState>,
533        status_tx: mpsc::UnboundedSender<String>,
534        server_id: crate::state::ServerId,
535        protocol: Arc<KafkaProtocol>,
536    ) -> Result<()> {
537        let (reader, write_half) = tokio::io::split(stream);
538        let write_half: SharedWrite = Arc::new(Mutex::new(write_half));
539
540        // Peer messaging, registered BEFORE the first read. Kafka is client-speaks-first and
541        // the broker says nothing until a request arrives, so a `*` manual rule parks the very
542        // first request this connection makes - and the operator being asked to decide about it
543        // must be able to reach, or hang up, the connection while it waits. Registering once
544        // the session was under way would leave exactly that case unreachable.
545        let peer_rx = crate::server::peer_support::register_peer_channel(
546            &app_state,
547            server_id,
548            connection_id.as_u32(),
549        )
550        .await;
551        crate::server::peer_support::spawn_peer_command_task(
552            peer_rx,
553            protocol.clone(),
554            app_state.clone(),
555            server_id,
556            connection_id.as_u32(),
557            write_half.clone(),
558            status_tx.clone(),
559        );
560
561        let result = Self::run_session(
562            reader,
563            &write_half,
564            peer_addr,
565            local_addr,
566            connection_id,
567            server,
568            llm_client,
569            app_state.clone(),
570            status_tx.clone(),
571            server_id,
572            protocol,
573        )
574        .await;
575
576        // Every exit path lands here - clean EOF, an invalid size prefix, an unparseable
577        // header, a write error, an injected disconnect. The `?`s all live in `run_session`,
578        // so none of them can skip this. Dropping the handle also ends the peer command task;
579        // the explicit shutdown makes the FIN immediate rather than waiting on it.
580        app_state
581            .remove_peer_handle(server_id, connection_id.as_u32())
582            .await;
583        let _ = write_half.lock().await.shutdown().await;
584        result
585    }
586
587    /// Read one length-prefixed request, answer it, repeat.
588    #[allow(clippy::too_many_arguments)]
589    async fn run_session(
590        mut reader: tokio::io::ReadHalf<TcpStream>,
591        write_half: &SharedWrite,
592        peer_addr: SocketAddr,
593        local_addr: SocketAddr,
594        connection_id: ConnectionId,
595        server: Arc<KafkaServer>,
596        llm_client: OllamaClient,
597        app_state: Arc<AppState>,
598        status_tx: mpsc::UnboundedSender<String>,
599        server_id: crate::state::ServerId,
600        protocol: Arc<KafkaProtocol>,
601    ) -> Result<()> {
602        let mut buffer = vec![0u8; 8192];
603        let mut answered_one = false;
604
605        loop {
606            // Both reads below are bounded by this deadline and nothing else is. The LLM
607            // round-trip that answers the request, and a `manual` rule parking it for a human
608            // (`src/state/intercepts.rs`, 300s by default), happen further down the loop body
609            // once the whole request has arrived — so neither sits inside a read deadline and
610            // neither can be cut short by one. What is bounded is a peer holding the connection
611            // while sending nothing, and a peer that announces a request size and then stalls
612            // part-way through delivering it.
613            let read_timeout = if answered_one {
614                IDLE_BETWEEN_REQUESTS_TIMEOUT
615            } else {
616                FIRST_REQUEST_READ_TIMEOUT
617            };
618
619            // Read the size prefix with read_exact. A plain read() may return 1-3
620            // bytes, and the old code parsed all four regardless, mixing in stale
621            // bytes from the previous message.
622            let mut size_prefix = [0u8; 4];
623            let prefix_read =
624                match tokio::time::timeout(read_timeout, reader.read_exact(&mut size_prefix)).await
625                {
626                    Ok(read) => read,
627                    Err(_) => {
628                        Log::new(Some(&status_tx)).debug(format!(
629                            "Kafka client {} sent nothing for {}s; closing idle connection",
630                            peer_addr,
631                            read_timeout.as_secs()
632                        ));
633                        break;
634                    }
635                };
636            match prefix_read {
637                Ok(_) => {}
638                Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
639                    Log::new(Some(&status_tx))
640                        .debug(format!("Kafka client {} disconnected", peer_addr));
641                    break;
642                }
643                Err(e) => return Err(e.into()),
644            }
645
646            // Validate the declared size before it reaches any allocator. This is
647            // unauthenticated input: `i32 as usize` sign-extends, so a prefix of
648            // 0x80000000 became ~1.8e19 and aborted the process on Vec::resize,
649            // while 0x7fffffff zeroed 2 GiB per connection. Sizes below the 8-byte
650            // request header are equally unusable. The comparison is done in i64 so
651            // that `declared` can never wrap.
652            let declared = i32::from_be_bytes(size_prefix);
653            if declared < MIN_REQUEST_BYTES || declared as i64 > MAX_REQUEST_BYTES as i64 {
654                Log::new(Some(&status_tx)).warn(format!(
655                    "Kafka client {} declared an invalid request size of {} bytes (allowed {}..={}); closing connection",
656                    peer_addr, declared, MIN_REQUEST_BYTES, MAX_REQUEST_BYTES
657                ));
658                break;
659            }
660            let message_size = declared as usize;
661
662            // Grow only: the buffer must never shrink below the size prefix.
663            if buffer.len() < message_size {
664                buffer.resize(message_size, 0);
665            }
666            match tokio::time::timeout(read_timeout, reader.read_exact(&mut buffer[..message_size]))
667                .await
668            {
669                Ok(read) => read?,
670                Err(_) => {
671                    Log::new(Some(&status_tx)).warn(format!(
672                        "Kafka client {} announced a {}-byte request and stalled for {}s \
673                         without delivering it; closing",
674                        peer_addr,
675                        message_size,
676                        read_timeout.as_secs()
677                    ));
678                    break;
679                }
680            };
681            // A whole request has arrived: from here on this peer is an established client and
682            // the longer idle bound applies.
683            answered_one = true;
684
685            app_state
686                .update_connection_stats(
687                    server_id,
688                    connection_id,
689                    Some((message_size + 4) as u64),
690                    None,
691                    Some(1),
692                    None,
693                )
694                .await;
695
696            let log = Log::new(Some(&status_tx));
697            log.debug(format!(
698                "Kafka received {} bytes from {}",
699                message_size, peer_addr
700            ));
701
702            // TRACE: hex dump, capped — hex::encode doubles the payload, so a
703            // maximum-size request would otherwise build a 200 MiB String.
704            let hex_len = message_size.min(MAX_TRACE_HEX_BYTES);
705            log.trace(format!(
706                "Kafka raw request: {} bytes, first {} hex: {}",
707                message_size,
708                hex_len,
709                hex::encode(&buffer[..hex_len])
710            ));
711
712            // api_key, api_version and correlation_id sit at fixed offsets in every
713            // request header version, so they can be read before choosing a header
714            // version to decode with. message_size >= MIN_REQUEST_BYTES (8) is already
715            // enforced above, so these slices are always in bounds.
716            let api_key_raw = i16::from_be_bytes([buffer[0], buffer[1]]);
717            let api_version = i16::from_be_bytes([buffer[2], buffer[3]]);
718            let correlation_id = i32::from_be_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]);
719
720            let api_key = match ApiKey::try_from(api_key_raw) {
721                Ok(k) => k,
722                Err(_) => {
723                    Log::new(Some(&status_tx)).error(format!(
724                        "Kafka client {} sent unknown API key {}; closing connection rather than \
725                         inventing a response body",
726                        peer_addr, api_key_raw
727                    ));
728                    break;
729                }
730            };
731
732            let supported = supported_versions(api_key_raw);
733            let version_ok =
734                matches!(supported, Some((min, max)) if api_version >= min && api_version <= max);
735
736            if !version_ok {
737                if api_key == ApiKey::ApiVersions {
738                    // Kafka's one special case: an ApiVersions request at a version the
739                    // broker does not implement is answered with UNSUPPORTED_VERSION and
740                    // the supported-API table, both encoded at v0, which every client
741                    // knows how to read. That is how clients negotiate downwards.
742                    Log::new(Some(&status_tx)).debug(format!(
743                        "Kafka ApiVersions v{} unsupported; replying at v0 with error {}",
744                        api_version, ERR_UNSUPPORTED_VERSION
745                    ));
746                    let body =
747                        Self::api_versions_response(0, correlation_id, ERR_UNSUPPORTED_VERSION)?;
748                    Self::write_frame(
749                        write_half,
750                        &body,
751                        &app_state,
752                        server_id,
753                        connection_id,
754                        &status_tx,
755                        peer_addr,
756                    )
757                    .await?;
758                    continue;
759                }
760                let range = supported
761                    .map(|(min, max)| format!("v{}..=v{}", min, max))
762                    .unwrap_or_else(|| "not implemented".to_string());
763                Log::new(Some(&status_tx)).error(format!(
764                    "Kafka client {} requested {:?} v{} which this broker does not implement ({}). \
765                     Closing the connection — a wrong-shaped body would be worse than an error.",
766                    peer_addr, api_key, api_version, range
767                ));
768                break;
769            }
770
771            // Decode the header at the version this (api_key, api_version) pair implies,
772            // then continue with the same cursor so the body starts where the header
773            // ended. Hardcoding 0 here left the cursor inside client_id for every real
774            // client and garbled every field after it.
775            let header_version = api_key.request_header_version(api_version);
776            let mut cursor = std::io::Cursor::new(&buffer[..message_size]);
777            let header = match RequestHeader::decode(&mut cursor, header_version) {
778                Ok(h) => h,
779                Err(e) => {
780                    Log::new(Some(&status_tx)).error(format!(
781                        "Kafka client {}: unparseable request header for {:?} v{} (header v{}): {}; closing",
782                        peer_addr, api_key, api_version, header_version, e
783                    ));
784                    break;
785                }
786            };
787
788            let client_id = header
789                .client_id
790                .as_ref()
791                .map(|c| c.to_string())
792                .unwrap_or_default();
793
794            Log::new(Some(&status_tx)).debug(format!(
795                "Kafka request: {:?} v{} correlation_id={} client_id={:?}",
796                api_key, api_version, header.correlation_id, client_id
797            ));
798
799            let ctx = RequestCtx {
800                api_version,
801                correlation_id: header.correlation_id,
802                client_id,
803                peer_addr,
804                local_addr,
805                connection_id,
806                server_id,
807            };
808
809            let response_bytes: Option<Vec<u8>> = match api_key {
810                ApiKey::ApiVersions => Some(Self::api_versions_response(
811                    api_version,
812                    header.correlation_id,
813                    ERR_NONE,
814                )?),
815                ApiKey::Metadata => Some(
816                    Self::handle_metadata(
817                        &mut cursor,
818                        &ctx,
819                        &server,
820                        &llm_client,
821                        &app_state,
822                        &status_tx,
823                        &protocol,
824                    )
825                    .await?,
826                ),
827                ApiKey::Produce => {
828                    Self::handle_produce(
829                        &mut cursor,
830                        &ctx,
831                        &llm_client,
832                        &app_state,
833                        &status_tx,
834                        &protocol,
835                    )
836                    .await?
837                }
838                ApiKey::Fetch => Some(
839                    Self::handle_fetch(
840                        &mut cursor,
841                        &ctx,
842                        &server,
843                        &llm_client,
844                        &app_state,
845                        &status_tx,
846                        &protocol,
847                    )
848                    .await?,
849                ),
850                ApiKey::OffsetCommit => Some(
851                    Self::handle_offset_commit(
852                        &mut cursor,
853                        &ctx,
854                        &llm_client,
855                        &app_state,
856                        &status_tx,
857                        &protocol,
858                    )
859                    .await?,
860                ),
861                // Unreachable: SUPPORTED_APIS gates the match above.
862                other => {
863                    Log::new(Some(&status_tx)).error(format!(
864                        "Kafka: API {:?} is advertised but has no handler; closing connection",
865                        other
866                    ));
867                    break;
868                }
869            };
870
871            match response_bytes {
872                Some(body) => {
873                    Self::write_frame(
874                        write_half,
875                        &body,
876                        &app_state,
877                        server_id,
878                        connection_id,
879                        &status_tx,
880                        peer_addr,
881                    )
882                    .await?;
883                }
884                None => {
885                    // acks=0 Produce: the producer expects no reply at all.
886                    Log::new(Some(&status_tx)).debug(format!(
887                        "Kafka: no response written for correlation_id={} (acks=0)",
888                        header.correlation_id
889                    ));
890                }
891            }
892        }
893
894        Ok(())
895    }
896
897    async fn write_frame(
898        write_half: &SharedWrite,
899        body: &[u8],
900        app_state: &Arc<AppState>,
901        server_id: crate::state::ServerId,
902        connection_id: ConnectionId,
903        status_tx: &mpsc::UnboundedSender<String>,
904        peer_addr: SocketAddr,
905    ) -> Result<()> {
906        let size = i32::try_from(body.len()).map_err(|_| {
907            anyhow::anyhow!(
908                "Kafka response of {} bytes exceeds the i32 size prefix",
909                body.len()
910            )
911        })?;
912        // One lock for the whole frame: the size prefix and the body are a single unit on the
913        // wire, and an injected write landing between them would desynchronise the client for
914        // the rest of the connection. The guard is dropped before the stats update so nothing
915        // awaits while holding the write half.
916        {
917            let mut writer = write_half.lock().await;
918            writer.write_all(&size.to_be_bytes()).await?;
919            writer.write_all(body).await?;
920            writer.flush().await?;
921        }
922
923        app_state
924            .update_connection_stats(
925                server_id,
926                connection_id,
927                None,
928                Some((body.len() + 4) as u64),
929                None,
930                Some(1),
931            )
932            .await;
933
934        let log = Log::new(Some(status_tx));
935        log.debug(format!("Kafka sent {} bytes to {}", body.len(), peer_addr));
936        let hex_len = body.len().min(MAX_TRACE_HEX_BYTES);
937        log.trace(format!(
938            "Kafka raw response: {} bytes, first {} hex: {}",
939            body.len(),
940            hex_len,
941            hex::encode(&body[..hex_len])
942        ));
943        Ok(())
944    }
945
946    /// ApiVersions is answered by Rust, not the model: it advertises what this *code*
947    /// can parse and encode, which is not a content decision.
948    fn api_versions_response(
949        api_version: i16,
950        correlation_id: i32,
951        error_code: i16,
952    ) -> Result<Vec<u8>> {
953        use kafka_protocol::messages::api_versions_response::ApiVersion;
954
955        let api_keys: Vec<ApiVersion> = SUPPORTED_APIS
956            .iter()
957            .map(|(key, min, max)| {
958                ApiVersion::default()
959                    .with_api_key(*key)
960                    .with_min_version(*min)
961                    .with_max_version(*max)
962            })
963            .collect();
964
965        let response = ApiVersionsResponse::default()
966            .with_error_code(error_code)
967            .with_api_keys(api_keys)
968            .with_throttle_time_ms(0);
969
970        let mut buf = Vec::new();
971        ResponseHeader::default()
972            .with_correlation_id(correlation_id)
973            .encode(
974                &mut buf,
975                ApiKey::ApiVersions.response_header_version(api_version),
976            )?;
977        response.encode(&mut buf, api_version)?;
978        Ok(buf)
979    }
980
981    /// Metadata: the model decides which topics exist and how they are partitioned.
982    /// Rust decides which broker address to advertise, because only Rust knows it.
983    #[allow(clippy::too_many_arguments)]
984    async fn handle_metadata(
985        cursor: &mut std::io::Cursor<&[u8]>,
986        ctx: &RequestCtx,
987        server: &Arc<KafkaServer>,
988        llm_client: &OllamaClient,
989        app_state: &Arc<AppState>,
990        status_tx: &mpsc::UnboundedSender<String>,
991        protocol: &Arc<KafkaProtocol>,
992    ) -> Result<Vec<u8>> {
993        use kafka_protocol::messages::metadata_response::{
994            MetadataResponseBroker, MetadataResponsePartition, MetadataResponseTopic,
995        };
996
997        let request = MetadataRequest::decode(cursor, ctx.api_version)?;
998
999        // `topics: None` means "all topics"; `Some([])` means "no topics". From v12 a topic
1000        // may be named by id alone; an id this broker never reported is UNKNOWN_TOPIC_ID.
1001        let mut unknown_ids = Vec::new();
1002        let requested_topics: Vec<String> = request
1003            .topics
1004            .as_ref()
1005            .map(|ts| {
1006                ts.iter()
1007                    .filter_map(|t| match &t.name {
1008                        Some(n) => Some(n.to_string()),
1009                        None => {
1010                            let name = server.topic_name(&t.topic_id);
1011                            if name.is_none() {
1012                                unknown_ids.push(t.topic_id);
1013                            }
1014                            name
1015                        }
1016                    })
1017                    .collect()
1018            })
1019            .unwrap_or_default();
1020        let wants_all_topics = request.topics.is_none();
1021
1022        Log::new(Some(status_tx)).debug(format!(
1023            "Kafka metadata request: all_topics={}, topics={:?}",
1024            wants_all_topics, requested_topics
1025        ));
1026
1027        let event = Event::new(
1028            &METADATA_REQUEST_EVENT,
1029            json!({
1030                "requested_topics": requested_topics,
1031                "all_topics": wants_all_topics,
1032                "client_id": ctx.client_id,
1033                "api_version": ctx.api_version,
1034            }),
1035        );
1036
1037        let reply = Self::ask_model(
1038            &event,
1039            "metadata_response",
1040            llm_client,
1041            app_state,
1042            status_tx,
1043            ctx,
1044            protocol,
1045        )
1046        .await;
1047
1048        // Whatever happens, the client needs a reachable broker address or it cannot
1049        // proceed. This is the one default in the module and it is transport truth.
1050        let default_broker = MetadataResponseBroker::default()
1051            .with_node_id(BrokerId(server.broker_id))
1052            .with_host(StrBytes::from_string(server.advertised_host.clone()))
1053            .with_port(ctx.local_addr.port() as i32);
1054
1055        let (brokers, mut response_topics) = match &reply {
1056            Reply::Data(data) => {
1057                let brokers = data
1058                    .get("brokers")
1059                    .and_then(|v| v.as_array())
1060                    .map(|list| {
1061                        list.iter()
1062                            .filter_map(|b| {
1063                                let host = b.get("host").and_then(|v| v.as_str())?;
1064                                let port = b.get("port").and_then(|v| v.as_i64())?;
1065                                if !(1..=65535).contains(&port) {
1066                                    warn!("Kafka metadata_response: broker port {} out of range, dropped", port);
1067                                    return None;
1068                                }
1069                                let id = b.get("id").and_then(|v| v.as_i64()).unwrap_or(server.broker_id as i64);
1070                                Some(
1071                                    MetadataResponseBroker::default()
1072                                        .with_node_id(BrokerId(id.clamp(0, i32::MAX as i64) as i32))
1073                                        .with_host(StrBytes::from_string(host.to_string()))
1074                                        .with_port(port as i32),
1075                                )
1076                            })
1077                            .collect::<Vec<_>>()
1078                    })
1079                    .filter(|v: &Vec<_>| !v.is_empty())
1080                    .unwrap_or_else(|| {
1081                        Log::new(Some(status_tx)).debug(format!(
1082                            "Kafka metadata_response named no usable broker; advertising this server ({}:{})",
1083                            server.advertised_host,
1084                            ctx.local_addr.port()
1085                        ));
1086                        vec![default_broker.clone()]
1087                    });
1088
1089                let mut topics = Vec::new();
1090                if let Some(list) = data.get("topics").and_then(|v| v.as_array()) {
1091                    for t in list {
1092                        let Some(name) = t.get("name").and_then(|v| v.as_str()) else {
1093                            Log::new(Some(status_tx)).warn(
1094                                "Kafka metadata_response: topic entry without a 'name' ignored",
1095                            );
1096                            continue;
1097                        };
1098                        // Kafka refuses such a name; a requested one falls through to the
1099                        // unknown-topic answer below rather than being described.
1100                        if name.len() > MAX_TOPIC_NAME_LEN {
1101                            Log::new(Some(status_tx)).warn(format!(
1102                                "Kafka metadata_response: topic name of {} bytes ignored; \
1103                                 Kafka's limit is {MAX_TOPIC_NAME_LEN}",
1104                                name.len()
1105                            ));
1106                            continue;
1107                        }
1108                        let topic_error =
1109                            clamp_error_code(t.get("error_code").and_then(|v| v.as_i64()));
1110
1111                        let mut partitions = Vec::new();
1112                        if let Some(plist) = t.get("partitions").and_then(|v| v.as_array()) {
1113                            for p in plist {
1114                                let idx = p.get("partition").and_then(|v| v.as_i64()).unwrap_or(0);
1115                                if !(0..=MAX_PARTITIONS as i64).contains(&idx) {
1116                                    Log::new(Some(status_tx)).warn(format!(
1117                                        "Kafka metadata_response: partition index {} for '{}' outside 0..={}, dropped",
1118                                        idx, name, MAX_PARTITIONS
1119                                    ));
1120                                    continue;
1121                                }
1122                                let leader = p
1123                                    .get("leader")
1124                                    .and_then(|v| v.as_i64())
1125                                    .unwrap_or(server.broker_id as i64)
1126                                    .clamp(0, i32::MAX as i64)
1127                                    as i32;
1128                                let replicas: Vec<BrokerId> = p
1129                                    .get("replicas")
1130                                    .and_then(|v| v.as_array())
1131                                    .map(|a| {
1132                                        a.iter()
1133                                            .filter_map(|x| x.as_i64())
1134                                            .map(|x| BrokerId(x.clamp(0, i32::MAX as i64) as i32))
1135                                            .collect()
1136                                    })
1137                                    .filter(|v: &Vec<BrokerId>| !v.is_empty())
1138                                    .unwrap_or_else(|| vec![BrokerId(leader)]);
1139                                partitions.push(
1140                                    MetadataResponsePartition::default()
1141                                        .with_partition_index(idx as i32)
1142                                        .with_leader_id(BrokerId(leader))
1143                                        .with_replica_nodes(replicas.clone())
1144                                        .with_isr_nodes(replicas)
1145                                        .with_error_code(ERR_NONE),
1146                                );
1147                            }
1148                        }
1149
1150                        if partitions.is_empty() && topic_error == ERR_NONE {
1151                            // A topic with no partitions is unusable to every client.
1152                            Log::new(Some(status_tx)).debug(format!(
1153                                "Kafka metadata_response: topic '{}' declared no partitions; assuming a single partition led by broker {}",
1154                                name, server.broker_id
1155                            ));
1156                            partitions.push(
1157                                MetadataResponsePartition::default()
1158                                    .with_partition_index(0)
1159                                    .with_leader_id(BrokerId(server.broker_id))
1160                                    .with_replica_nodes(vec![BrokerId(server.broker_id)])
1161                                    .with_isr_nodes(vec![BrokerId(server.broker_id)])
1162                                    .with_error_code(ERR_NONE),
1163                            );
1164                        }
1165
1166                        let mut topic = MetadataResponseTopic::default()
1167                            .with_name(Some(StrBytes::from_string(name.to_string()).into()))
1168                            .with_partitions(partitions)
1169                            .with_error_code(topic_error);
1170                        if ctx.api_version >= 10 {
1171                            if let Some(id) = server.report_topic(name) {
1172                                topic = topic.with_topic_id(id);
1173                            }
1174                        }
1175                        topics.push(topic);
1176                    }
1177                }
1178                (brokers, topics)
1179            }
1180            Reply::Error(code) => {
1181                Log::new(Some(status_tx)).debug(format!(
1182                    "Kafka metadata: model refused with error code {}",
1183                    code
1184                ));
1185                (vec![default_broker.clone()], Vec::new())
1186            }
1187            Reply::NoAnswer(_) => (vec![default_broker.clone()], Vec::new()),
1188        };
1189
1190        // Any topic the client explicitly asked about but the model did not describe is
1191        // reported as unknown (or with the model's chosen error code). Silence about a
1192        // requested topic must never look like success.
1193        let fallback_error = match &reply {
1194            Reply::Data(_) => ERR_UNKNOWN_TOPIC_OR_PARTITION,
1195            Reply::Error(code) => *code,
1196            Reply::NoAnswer(code) => *code,
1197        };
1198        for name in &requested_topics {
1199            let described = response_topics
1200                .iter()
1201                .any(|t| t.name.as_ref().map(|n| n.to_string()).as_deref() == Some(name.as_str()));
1202            if !described {
1203                response_topics.push(
1204                    MetadataResponseTopic::default()
1205                        .with_name(Some(StrBytes::from_string(name.clone()).into()))
1206                        .with_error_code(fallback_error),
1207                );
1208            }
1209        }
1210        for id in unknown_ids {
1211            response_topics.push(
1212                MetadataResponseTopic::default()
1213                    .with_name(None)
1214                    .with_topic_id(id)
1215                    .with_error_code(ERR_UNKNOWN_TOPIC_ID),
1216            );
1217        }
1218
1219        if matches!(reply, Reply::NoAnswer(_)) && requested_topics.is_empty() {
1220            Log::new(Some(status_tx)).warn(
1221                "Kafka metadata: no answer from the model and no specific topic was requested; \
1222                 replying with the broker list and an empty topic list",
1223            );
1224        }
1225
1226        Log::new(Some(status_tx)).info(format!(
1227            "Kafka metadata reply: {} broker(s), {} topic(s)",
1228            brokers.len(),
1229            response_topics.len()
1230        ));
1231
1232        let response = MetadataResponse::default()
1233            .with_brokers(brokers)
1234            .with_cluster_id(Some(StrBytes::from_string(server.cluster_id.clone())))
1235            .with_controller_id(BrokerId(server.broker_id))
1236            .with_topics(response_topics)
1237            .with_throttle_time_ms(0);
1238
1239        let mut buf = Vec::new();
1240        ResponseHeader::default()
1241            .with_correlation_id(ctx.correlation_id)
1242            .encode(
1243                &mut buf,
1244                ApiKey::Metadata.response_header_version(ctx.api_version),
1245            )?;
1246        response.encode(&mut buf, ctx.api_version)?;
1247        Ok(buf)
1248    }
1249
1250    /// Produce: Rust decodes the record batch, the model decides whether it is accepted
1251    /// and at what offset. Nothing is stored.
1252    ///
1253    /// Returns `None` when `acks == 0`, because the producer is not waiting for a reply.
1254    #[allow(clippy::too_many_arguments)]
1255    async fn handle_produce(
1256        cursor: &mut std::io::Cursor<&[u8]>,
1257        ctx: &RequestCtx,
1258        llm_client: &OllamaClient,
1259        app_state: &Arc<AppState>,
1260        status_tx: &mpsc::UnboundedSender<String>,
1261        protocol: &Arc<KafkaProtocol>,
1262    ) -> Result<Option<Vec<u8>>> {
1263        use kafka_protocol::messages::produce_response::{
1264            PartitionProduceResponse, TopicProduceResponse,
1265        };
1266
1267        let request = ProduceRequest::decode(cursor, ctx.api_version)?;
1268        let acks = request.acks;
1269
1270        let mut topic_responses = Vec::new();
1271        let mut units = 0usize;
1272
1273        for topic_data in &request.topic_data {
1274            let topic_name = topic_data.name.to_string();
1275            let mut partition_responses = Vec::new();
1276
1277            for partition_data in &topic_data.partition_data {
1278                let partition_idx = partition_data.index;
1279
1280                if !(0..=MAX_PARTITIONS).contains(&partition_idx) {
1281                    Log::new(Some(status_tx)).warn(format!(
1282                        "Kafka produce for '{}' names out-of-range partition {} (0..={}); rejected",
1283                        topic_name, partition_idx, MAX_PARTITIONS
1284                    ));
1285                    partition_responses.push(
1286                        PartitionProduceResponse::default()
1287                            .with_index(partition_idx)
1288                            .with_error_code(ERR_UNKNOWN_TOPIC_OR_PARTITION)
1289                            .with_base_offset(-1),
1290                    );
1291                    continue;
1292                }
1293
1294                units += 1;
1295                if units > MAX_UNITS_PER_REQUEST {
1296                    Log::new(Some(status_tx)).warn(format!(
1297                        "Kafka produce names more than {} topic-partitions in one request; \
1298                         remaining ones rejected without consulting the model",
1299                        MAX_UNITS_PER_REQUEST
1300                    ));
1301                    partition_responses.push(
1302                        PartitionProduceResponse::default()
1303                            .with_index(partition_idx)
1304                            .with_error_code(ERR_UNKNOWN_SERVER_ERROR)
1305                            .with_base_offset(-1),
1306                    );
1307                    continue;
1308                }
1309
1310                // Decode the batch. A batch we cannot parse is CORRUPT_MESSAGE — it used
1311                // to be replaced with an empty placeholder record and acknowledged as
1312                // success, which told the producer its data was durable when it had been
1313                // silently dropped.
1314                let records = match &partition_data.records {
1315                    Some(raw) => {
1316                        let owned = Bytes::copy_from_slice(raw.as_ref());
1317                        let mut rc = std::io::Cursor::new(owned);
1318                        match RecordBatchDecoder::decode_with_custom_compression::<
1319                            _,
1320                            fn(&mut Bytes, Compression) -> Result<std::io::Cursor<Bytes>>,
1321                        >(
1322                            &mut rc,
1323                            None::<fn(&mut Bytes, Compression) -> Result<std::io::Cursor<Bytes>>>,
1324                        ) {
1325                            Ok(r) => r,
1326                            Err(e) => {
1327                                Log::new(Some(status_tx)).warn(format!(
1328                                    "Kafka produce to '{}' partition {}: unparseable record batch ({}); \
1329                                     replying CORRUPT_MESSAGE",
1330                                    topic_name, partition_idx, e
1331                                ));
1332                                partition_responses.push(
1333                                    PartitionProduceResponse::default()
1334                                        .with_index(partition_idx)
1335                                        .with_error_code(ERR_CORRUPT_MESSAGE)
1336                                        .with_base_offset(-1),
1337                                );
1338                                continue;
1339                            }
1340                        }
1341                    }
1342                    None => Vec::new(),
1343                };
1344
1345                let event_records: Vec<Value> = records
1346                    .iter()
1347                    .take(MAX_EVENT_RECORDS)
1348                    .map(|r| {
1349                        let (key, key_encoding) = r
1350                            .key
1351                            .as_ref()
1352                            .map(|k| encode_field(k, MAX_EVENT_VALUE_BYTES))
1353                            .unwrap_or((Value::Null, "utf8"));
1354                        let (value, value_encoding) = r
1355                            .value
1356                            .as_ref()
1357                            .map(|v| encode_field(v, MAX_EVENT_VALUE_BYTES))
1358                            .unwrap_or((Value::Null, "utf8"));
1359                        json!({
1360                            "offset": r.offset,
1361                            "timestamp": r.timestamp,
1362                            "key": key,
1363                            "key_encoding": key_encoding,
1364                            "value": value,
1365                            "value_encoding": value_encoding,
1366                        })
1367                    })
1368                    .collect();
1369
1370                let first_key = records
1371                    .first()
1372                    .and_then(|r| r.key.as_ref())
1373                    .map(|k| encode_field(k, MAX_EVENT_VALUE_BYTES).0)
1374                    .unwrap_or(Value::Null);
1375                let first_value = records
1376                    .first()
1377                    .and_then(|r| r.value.as_ref())
1378                    .map(|v| encode_field(v, MAX_EVENT_VALUE_BYTES).0)
1379                    .unwrap_or(Value::String(String::new()));
1380
1381                let event = Event::new(
1382                    &PRODUCE_REQUEST_EVENT,
1383                    json!({
1384                        "topic": topic_name,
1385                        "partition": partition_idx,
1386                        "record_count": records.len(),
1387                        "first_key": first_key,
1388                        "first_value_preview": first_value,
1389                        "records": event_records,
1390                        "acks": acks,
1391                        "client_id": ctx.client_id,
1392                        "api_version": ctx.api_version,
1393                    }),
1394                );
1395
1396                let reply = Self::ask_model(
1397                    &event,
1398                    "produce_response",
1399                    llm_client,
1400                    app_state,
1401                    status_tx,
1402                    ctx,
1403                    protocol,
1404                )
1405                .await;
1406
1407                let (error_code, base_offset) = match reply {
1408                    Reply::Data(data) => {
1409                        let code =
1410                            clamp_error_code(data.get("error_code").and_then(|v| v.as_i64()));
1411                        let offset = data
1412                            .get("offset")
1413                            .and_then(|v| v.as_i64())
1414                            .unwrap_or_else(|| {
1415                                if code == ERR_NONE {
1416                                    Log::new(Some(status_tx)).warn(format!(
1417                                        "Kafka produce_response for '{}' partition {} omitted 'offset'; using 0",
1418                                        topic_name, partition_idx
1419                                    ));
1420                                }
1421                                0
1422                            });
1423                        if code == ERR_NONE {
1424                            Log::new(Some(status_tx)).info(format!(
1425                                "Kafka produce accepted: '{}' partition {} at offset {} ({} record(s))",
1426                                topic_name, partition_idx, offset, records.len()
1427                            ));
1428                            (code, offset.max(0))
1429                        } else {
1430                            (code, -1)
1431                        }
1432                    }
1433                    Reply::Error(code) => (code, -1),
1434                    Reply::NoAnswer(code) => (code, -1),
1435                };
1436
1437                partition_responses.push(
1438                    PartitionProduceResponse::default()
1439                        .with_index(partition_idx)
1440                        .with_base_offset(base_offset)
1441                        .with_error_code(error_code),
1442                );
1443            }
1444
1445            topic_responses.push(
1446                TopicProduceResponse::default()
1447                    .with_name(StrBytes::from_string(topic_name).into())
1448                    .with_partition_responses(partition_responses),
1449            );
1450        }
1451
1452        if acks == 0 {
1453            // Fire-and-forget producer: sending a response here would desynchronise it.
1454            return Ok(None);
1455        }
1456
1457        let response = ProduceResponse::default()
1458            .with_responses(topic_responses)
1459            .with_throttle_time_ms(0);
1460
1461        let mut buf = Vec::new();
1462        ResponseHeader::default()
1463            .with_correlation_id(ctx.correlation_id)
1464            .encode(
1465                &mut buf,
1466                ApiKey::Produce.response_header_version(ctx.api_version),
1467            )?;
1468        response.encode(&mut buf, ctx.api_version)?;
1469        Ok(Some(buf))
1470    }
1471
1472    /// Fetch: the model supplies the records. There is no log to read from.
1473    #[allow(clippy::too_many_arguments)]
1474    async fn handle_fetch(
1475        cursor: &mut std::io::Cursor<&[u8]>,
1476        ctx: &RequestCtx,
1477        server: &Arc<KafkaServer>,
1478        llm_client: &OllamaClient,
1479        app_state: &Arc<AppState>,
1480        status_tx: &mpsc::UnboundedSender<String>,
1481        protocol: &Arc<KafkaProtocol>,
1482    ) -> Result<Vec<u8>> {
1483        use kafka_protocol::messages::fetch_response::{FetchableTopicResponse, PartitionData};
1484
1485        let request = FetchRequest::decode(cursor, ctx.api_version)?;
1486        let request_max_bytes = request.max_bytes;
1487
1488        let mut topic_responses = Vec::new();
1489        let mut units = 0usize;
1490
1491        for topic in &request.topics {
1492            // From v13 a fetch names its topic by id; one this broker never reported cannot
1493            // be asked about, and every partition of it is UNKNOWN_TOPIC_ID.
1494            let by_id = ctx.api_version >= 13;
1495            let topic_name = if by_id {
1496                match server.topic_name(&topic.topic_id) {
1497                    Some(name) => name,
1498                    None => {
1499                        topic_responses.push(
1500                            FetchableTopicResponse::default()
1501                                .with_topic_id(topic.topic_id)
1502                                .with_partitions(
1503                                    topic
1504                                        .partitions
1505                                        .iter()
1506                                        .map(|p| {
1507                                            PartitionData::default()
1508                                                .with_partition_index(p.partition)
1509                                                .with_error_code(ERR_UNKNOWN_TOPIC_ID)
1510                                                .with_records(Some(Bytes::new()))
1511                                        })
1512                                        .collect(),
1513                                ),
1514                        );
1515                        continue;
1516                    }
1517                }
1518            } else {
1519                topic.topic.to_string()
1520            };
1521            let mut partition_responses = Vec::new();
1522
1523            for partition in &topic.partitions {
1524                let partition_idx = partition.partition;
1525                let fetch_offset = partition.fetch_offset;
1526
1527                if !(0..=MAX_PARTITIONS).contains(&partition_idx) {
1528                    partition_responses.push(
1529                        PartitionData::default()
1530                            .with_partition_index(partition_idx)
1531                            .with_error_code(ERR_UNKNOWN_TOPIC_OR_PARTITION)
1532                            .with_records(Some(Bytes::new())),
1533                    );
1534                    continue;
1535                }
1536
1537                units += 1;
1538                if units > MAX_UNITS_PER_REQUEST {
1539                    Log::new(Some(status_tx)).warn(format!(
1540                        "Kafka fetch names more than {} topic-partitions in one request; \
1541                         remaining ones rejected without consulting the model",
1542                        MAX_UNITS_PER_REQUEST
1543                    ));
1544                    partition_responses.push(
1545                        PartitionData::default()
1546                            .with_partition_index(partition_idx)
1547                            .with_error_code(ERR_UNKNOWN_SERVER_ERROR)
1548                            .with_records(Some(Bytes::new())),
1549                    );
1550                    continue;
1551                }
1552
1553                let max_bytes = if partition.partition_max_bytes > 0 {
1554                    partition.partition_max_bytes
1555                } else {
1556                    request_max_bytes
1557                };
1558
1559                let event = Event::new(
1560                    &FETCH_REQUEST_EVENT,
1561                    json!({
1562                        "topic": topic_name,
1563                        "partition": partition_idx,
1564                        "fetch_offset": fetch_offset,
1565                        "max_bytes": max_bytes,
1566                        "client_id": ctx.client_id,
1567                        "api_version": ctx.api_version,
1568                    }),
1569                );
1570
1571                let reply = Self::ask_model(
1572                    &event,
1573                    "fetch_response",
1574                    llm_client,
1575                    app_state,
1576                    status_tx,
1577                    ctx,
1578                    protocol,
1579                )
1580                .await;
1581
1582                let pd = match reply {
1583                    Reply::Data(data) => {
1584                        let list = data
1585                            .get("records")
1586                            .and_then(|v| v.as_array())
1587                            .cloned()
1588                            .unwrap_or_default();
1589                        match encode_record_batch(&list, fetch_offset, status_tx) {
1590                            Ok((bytes, high_watermark)) => PartitionData::default()
1591                                .with_partition_index(partition_idx)
1592                                .with_high_watermark(high_watermark)
1593                                .with_last_stable_offset(high_watermark)
1594                                .with_log_start_offset(0)
1595                                .with_records(Some(bytes))
1596                                .with_error_code(ERR_NONE),
1597                            Err(e) => {
1598                                Log::new(Some(status_tx)).warn(format!(
1599                                    "Kafka fetch for '{}' partition {}: the model's records could not be \
1600                                     encoded ({}); replying UNKNOWN_SERVER_ERROR",
1601                                    topic_name, partition_idx, e
1602                                ));
1603                                PartitionData::default()
1604                                    .with_partition_index(partition_idx)
1605                                    .with_error_code(ERR_UNKNOWN_SERVER_ERROR)
1606                                    .with_records(Some(Bytes::new()))
1607                            }
1608                        }
1609                    }
1610                    Reply::Error(code) => PartitionData::default()
1611                        .with_partition_index(partition_idx)
1612                        .with_error_code(code)
1613                        .with_records(Some(Bytes::new())),
1614                    Reply::NoAnswer(code) => PartitionData::default()
1615                        .with_partition_index(partition_idx)
1616                        .with_error_code(code)
1617                        .with_records(Some(Bytes::new())),
1618                };
1619
1620                partition_responses.push(pd);
1621            }
1622
1623            let response = FetchableTopicResponse::default().with_partitions(partition_responses);
1624            topic_responses.push(if by_id {
1625                response.with_topic_id(topic.topic_id)
1626            } else {
1627                response.with_topic(StrBytes::from_string(topic_name).into())
1628            });
1629        }
1630
1631        let response = FetchResponse::default()
1632            .with_responses(topic_responses)
1633            .with_throttle_time_ms(0);
1634
1635        let mut buf = Vec::new();
1636        ResponseHeader::default()
1637            .with_correlation_id(ctx.correlation_id)
1638            .encode(
1639                &mut buf,
1640                ApiKey::Fetch.response_header_version(ctx.api_version),
1641            )?;
1642        response.encode(&mut buf, ctx.api_version)?;
1643        Ok(buf)
1644    }
1645
1646    /// OffsetCommit: the model decides whether the commit is accepted. Nothing is stored,
1647    /// so a later OffsetFetch would not see it — and OffsetFetch is not implemented.
1648    #[allow(clippy::too_many_arguments)]
1649    async fn handle_offset_commit(
1650        cursor: &mut std::io::Cursor<&[u8]>,
1651        ctx: &RequestCtx,
1652        llm_client: &OllamaClient,
1653        app_state: &Arc<AppState>,
1654        status_tx: &mpsc::UnboundedSender<String>,
1655        protocol: &Arc<KafkaProtocol>,
1656    ) -> Result<Vec<u8>> {
1657        use kafka_protocol::messages::offset_commit_response::{
1658            OffsetCommitResponsePartition, OffsetCommitResponseTopic,
1659        };
1660
1661        let request = OffsetCommitRequest::decode(cursor, ctx.api_version)?;
1662        let group_id = request.group_id.to_string();
1663
1664        let mut topic_responses = Vec::new();
1665        let mut units = 0usize;
1666
1667        for topic in &request.topics {
1668            let topic_name = topic.name.to_string();
1669            let mut partition_responses = Vec::new();
1670
1671            for partition in &topic.partitions {
1672                let partition_idx = partition.partition_index;
1673
1674                units += 1;
1675                if units > MAX_UNITS_PER_REQUEST {
1676                    Log::new(Some(status_tx)).warn(format!(
1677                        "Kafka offset commit names more than {} topic-partitions in one request; \
1678                         remaining ones rejected without consulting the model",
1679                        MAX_UNITS_PER_REQUEST
1680                    ));
1681                    partition_responses.push(
1682                        OffsetCommitResponsePartition::default()
1683                            .with_partition_index(partition_idx)
1684                            .with_error_code(ERR_UNKNOWN_SERVER_ERROR),
1685                    );
1686                    continue;
1687                }
1688
1689                let event = Event::new(
1690                    &OFFSET_COMMIT_REQUEST_EVENT,
1691                    json!({
1692                        "group_id": group_id,
1693                        "topic": topic_name,
1694                        "partition": partition_idx,
1695                        "offset": partition.committed_offset,
1696                        "client_id": ctx.client_id,
1697                        "api_version": ctx.api_version,
1698                    }),
1699                );
1700
1701                let reply = Self::ask_model(
1702                    &event,
1703                    "offset_commit_response",
1704                    llm_client,
1705                    app_state,
1706                    status_tx,
1707                    ctx,
1708                    protocol,
1709                )
1710                .await;
1711
1712                let error_code = match reply {
1713                    Reply::Data(data) => {
1714                        clamp_error_code(data.get("error_code").and_then(|v| v.as_i64()))
1715                    }
1716                    Reply::Error(code) => code,
1717                    Reply::NoAnswer(code) => code,
1718                };
1719
1720                partition_responses.push(
1721                    OffsetCommitResponsePartition::default()
1722                        .with_partition_index(partition_idx)
1723                        .with_error_code(error_code),
1724                );
1725            }
1726
1727            topic_responses.push(
1728                OffsetCommitResponseTopic::default()
1729                    .with_name(StrBytes::from_string(topic_name).into())
1730                    .with_partitions(partition_responses),
1731            );
1732        }
1733
1734        let response = OffsetCommitResponse::default()
1735            .with_topics(topic_responses)
1736            .with_throttle_time_ms(0);
1737
1738        let mut buf = Vec::new();
1739        ResponseHeader::default()
1740            .with_correlation_id(ctx.correlation_id)
1741            .encode(
1742                &mut buf,
1743                ApiKey::OffsetCommit.response_header_version(ctx.api_version),
1744            )?;
1745        response.encode(&mut buf, ctx.api_version)?;
1746        Ok(buf)
1747    }
1748
1749    /// Raise one event and reduce the returned actions to a single decision.
1750    ///
1751    /// `call_llm` is the entry point for all three handling modes, so a script or static
1752    /// `event_handler` short-circuits the model here exactly as it does for every other
1753    /// protocol.
1754    async fn ask_model(
1755        event: &Event,
1756        expected_action: &str,
1757        llm_client: &OllamaClient,
1758        app_state: &Arc<AppState>,
1759        status_tx: &mpsc::UnboundedSender<String>,
1760        ctx: &RequestCtx,
1761        protocol: &Arc<KafkaProtocol>,
1762    ) -> Reply {
1763        let result = call_llm(
1764            llm_client,
1765            app_state,
1766            ctx.server_id,
1767            Some(ctx.connection_id),
1768            event,
1769            protocol.as_ref(),
1770        )
1771        .await;
1772
1773        let execution = match result {
1774            Ok(r) => r,
1775            Err(e) => {
1776                // The peer gets a category, the log gets the error - and the two ways of
1777                // failing are tagged distinctly, as `src/server/radius/` does, because on
1778                // the wire Kafka can only carry an error code and the log is the only place
1779                // the difference survives.
1780                let failure = crate::utils::WireFailure::classify(&e);
1781                let (code, class) = if failure.is_overloaded() {
1782                    (ERR_REQUEST_TIMED_OUT, "overloaded")
1783                } else {
1784                    (ERR_UNKNOWN_SERVER_ERROR, "unavailable")
1785                };
1786                Log::new(Some(status_tx)).warn(format!(
1787                    "Kafka: handler failed for {}: decision=fail_closed_llm_error class={} \
1788                     error_code={} error={}",
1789                    event.id(),
1790                    class,
1791                    code,
1792                    e
1793                ));
1794                return Reply::NoAnswer(code);
1795            }
1796        };
1797
1798        // First recognised action wins, so the model's ordering is respected and an
1799        // explicit refusal is never overridden by a later success.
1800        for result in execution.protocol_results {
1801            if let ActionResult::Custom { name, data } = result {
1802                if name == expected_action {
1803                    return Reply::Data(data);
1804                }
1805                // Falls through to `error_response` below; keeping the two apart in the log
1806                // is what makes a model's deliberate refusal distinguishable from silence.
1807                if name == "error_response" {
1808                    let mut code =
1809                        clamp_error_code(data.get("error_code").and_then(|v| v.as_i64()));
1810                    if code == ERR_NONE {
1811                        // "error_response" with error_code 0 is a contradiction; treating
1812                        // it as success would make a refusal indistinguishable from an ack.
1813                        Log::new(Some(status_tx)).warn(format!(
1814                            "Kafka: error_response for {} carried error_code 0; using \
1815                             UNKNOWN_SERVER_ERROR instead",
1816                            event.id()
1817                        ));
1818                        code = ERR_UNKNOWN_SERVER_ERROR;
1819                    }
1820                    Log::new(Some(status_tx)).info(format!(
1821                        "Kafka: {} refused: decision=model_reject error_code={}",
1822                        event.id(),
1823                        code
1824                    ));
1825                    return Reply::Error(code);
1826                }
1827                debug!(
1828                    "Kafka: ignoring action '{}' for {} (expected '{}' or 'error_response')",
1829                    name,
1830                    event.id(),
1831                    expected_action
1832                );
1833            }
1834        }
1835
1836        // The handler ran and produced nothing this event can use. Distinct from the
1837        // backend failing above: permanent as far as the client is concerned, and tagged
1838        // so an operator can tell the two apart in the log.
1839        Log::new(Some(status_tx)).warn(format!(
1840            "Kafka: {} produced no '{}' and no 'error_response': decision=model_silent \
1841             error_code={}",
1842            event.id(),
1843            expected_action,
1844            ERR_UNKNOWN_SERVER_ERROR
1845        ));
1846        Reply::NoAnswer(ERR_UNKNOWN_SERVER_ERROR)
1847    }
1848}
1849
1850/// Everything a handler needs about the request it is answering, so the handler
1851/// signatures stay readable.
1852struct RequestCtx {
1853    api_version: i16,
1854    correlation_id: i32,
1855    client_id: String,
1856    #[allow(dead_code)]
1857    peer_addr: SocketAddr,
1858    local_addr: SocketAddr,
1859    connection_id: ConnectionId,
1860    server_id: crate::state::ServerId,
1861}
1862
1863/// Clamp a model-supplied error code into i16, defaulting to success when absent.
1864fn clamp_error_code(v: Option<i64>) -> i16 {
1865    match v {
1866        Some(c) if c >= i16::MIN as i64 && c <= i16::MAX as i64 => c as i16,
1867        Some(_) => ERR_UNKNOWN_SERVER_ERROR,
1868        None => ERR_NONE,
1869    }
1870}
1871
1872/// Render record bytes for an event, as text when that is faithful and hex otherwise.
1873///
1874/// Models cannot read base64 and cannot be trusted to notice that a "string" is really
1875/// binary, so the encoding is stated explicitly alongside the value — the same shape the
1876/// TCP protocol settled on.
1877///
1878/// Shared with the Kafka *client* (`src/client/kafka/`) so that "hex means hex" is decided
1879/// once for both directions of the same connection.
1880pub(crate) fn encode_field(bytes: &[u8], max_len: usize) -> (Value, &'static str) {
1881    let printable = std::str::from_utf8(bytes).ok().filter(|s| {
1882        !s.chars()
1883            .any(|c| c.is_control() && c != '\n' && c != '\r' && c != '\t')
1884    });
1885    match printable {
1886        Some(s) => {
1887            let truncated = crate::utils::truncate::truncate_for_llm(s, max_len);
1888            (Value::String(truncated), "utf8")
1889        }
1890        None => {
1891            let take = bytes.len().min(max_len / 2);
1892            (Value::String(hex::encode(&bytes[..take])), "hex")
1893        }
1894    }
1895}
1896
1897/// Turn one action field into record bytes, honouring the declared encoding.
1898///
1899/// If a field is documented as accepting hex, the executor has to actually decode it —
1900/// anything else puts literal ASCII on the wire.
1901///
1902/// Shared with the Kafka *client* (`src/client/kafka/`), which has the same obligation on its
1903/// produce path.
1904pub(crate) fn decode_field(value: Option<&Value>, encoding: &str) -> Result<Option<Vec<u8>>> {
1905    let Some(value) = value else {
1906        return Ok(None);
1907    };
1908    if value.is_null() {
1909        return Ok(None);
1910    }
1911    let text = match value {
1912        Value::String(s) => s.clone(),
1913        other => other.to_string(),
1914    };
1915    match encoding {
1916        "hex" => {
1917            let cleaned: String = text.chars().filter(|c| !c.is_whitespace()).collect();
1918            let bytes = hex::decode(&cleaned)
1919                .map_err(|e| anyhow::anyhow!("invalid hex in record field: {}", e))?;
1920            Ok(Some(bytes))
1921        }
1922        "utf8" => Ok(Some(text.into_bytes())),
1923        other => Err(anyhow::anyhow!(
1924            "unknown record encoding '{}' (expected 'utf8' or 'hex')",
1925            other
1926        )),
1927    }
1928}
1929
1930/// Encode the model's records into a v2 record batch.
1931///
1932/// Returns the encoded batch and the high watermark the client should be told about
1933/// (one past the last offset in the batch).
1934///
1935/// Offsets are made contiguous from the batch's base. Kafka's v2 batch stores per-record
1936/// offset *deltas* from the base and the encoder requires `offset - sequence` to be
1937/// constant across a batch, so honouring arbitrary gaps would silently split or corrupt
1938/// the batch. A model-supplied offset below `fetch_offset` would be discarded by the
1939/// consumer, so it is raised to `fetch_offset` and logged.
1940fn encode_record_batch(
1941    records: &[Value],
1942    fetch_offset: i64,
1943    status_tx: &mpsc::UnboundedSender<String>,
1944) -> Result<(Bytes, i64)> {
1945    if records.is_empty() {
1946        return Ok((Bytes::new(), fetch_offset.max(0)));
1947    }
1948
1949    if records.len() > MAX_RECORDS_PER_FETCH {
1950        Log::new(Some(status_tx)).warn(format!(
1951            "Kafka fetch_response contained {} records; only the first {} are returned",
1952            records.len(),
1953            MAX_RECORDS_PER_FETCH
1954        ));
1955    }
1956
1957    let base_offset = records
1958        .first()
1959        .and_then(|r| r.get("offset"))
1960        .and_then(|v| v.as_i64())
1961        .unwrap_or(fetch_offset)
1962        .max(fetch_offset.max(0));
1963
1964    let now = chrono::Utc::now().timestamp_millis();
1965    let mut total_payload = 0usize;
1966    let mut out = Vec::new();
1967
1968    for (idx, r) in records.iter().take(MAX_RECORDS_PER_FETCH).enumerate() {
1969        let key_encoding = r
1970            .get("key_encoding")
1971            .and_then(|v| v.as_str())
1972            .or_else(|| r.get("encoding").and_then(|v| v.as_str()))
1973            .unwrap_or("utf8");
1974        let value_encoding = r
1975            .get("value_encoding")
1976            .and_then(|v| v.as_str())
1977            .or_else(|| r.get("encoding").and_then(|v| v.as_str()))
1978            .unwrap_or("utf8");
1979
1980        let key = decode_field(r.get("key"), key_encoding)?;
1981        let value = decode_field(r.get("value"), value_encoding)?;
1982
1983        total_payload += key.as_ref().map(|k| k.len()).unwrap_or(0)
1984            + value.as_ref().map(|v| v.len()).unwrap_or(0);
1985        if total_payload > MAX_FETCH_PAYLOAD_BYTES {
1986            Log::new(Some(status_tx)).warn(format!(
1987                "Kafka fetch_response payload exceeded {} bytes; truncating the batch at {} record(s)",
1988                MAX_FETCH_PAYLOAD_BYTES, idx
1989            ));
1990            break;
1991        }
1992
1993        let declared = r.get("offset").and_then(|v| v.as_i64());
1994        let offset = base_offset + idx as i64;
1995        if let Some(d) = declared {
1996            if d != offset {
1997                debug!(
1998                    "Kafka fetch_response: record offset {} rewritten to {} to keep the batch contiguous",
1999                    d, offset
2000                );
2001            }
2002        }
2003
2004        out.push(Record {
2005            transactional: false,
2006            control: false,
2007            partition_leader_epoch: 0,
2008            producer_id: -1,
2009            producer_epoch: -1,
2010            timestamp_type: TimestampType::Creation,
2011            offset,
2012            sequence: idx as i32,
2013            timestamp: r.get("timestamp").and_then(|v| v.as_i64()).unwrap_or(now),
2014            key: key.map(Bytes::from),
2015            value: value.map(Bytes::from),
2016            headers: Default::default(),
2017        });
2018    }
2019
2020    if out.is_empty() {
2021        return Ok((Bytes::new(), base_offset));
2022    }
2023
2024    let mut buf = Vec::new();
2025    let options = RecordEncodeOptions {
2026        version: 2,
2027        compression: Compression::None,
2028    };
2029    RecordBatchEncoder::encode_with_custom_compression::<
2030        _,
2031        _,
2032        fn(&mut bytes::BytesMut, &mut Vec<u8>, Compression) -> Result<()>,
2033    >(&mut buf, &out, &options, None)?;
2034
2035    let high_watermark = base_offset + out.len() as i64;
2036    Ok((Bytes::from(buf), high_watermark))
2037}