Skip to main content

netget/client/kafka/
mod.rs

1//! Kafka client — pure Rust, no librdkafka.
2//!
3//! # Wire format: shared with the broker, not reimplemented
4//!
5//! Every byte is produced and consumed by `kafka-protocol`'s code-generated codecs, reached
6//! through the `pub use kafka_protocol` re-export in [`crate::server::kafka`]. Both halves
7//! live behind the same `kafka` Cargo feature, so there is nothing extra to gate and no
8//! second copy of the schemas to keep in sync — the same arrangement the BGP client has with
9//! `src/server/bgp/wire.rs`.
10//!
11//! This client uses the codecs in the *client* direction: it encodes requests and decodes
12//! responses, where the broker decodes requests and encodes responses. Nothing here calls a
13//! function `src/server/kafka/mod.rs` wrote, apart from the two record-field encoding helpers
14//! that decide how bytes are shown to (and taken from) a model.
15//!
16//! ## What this replaced
17//!
18//! An `rdkafka` (librdkafka) client that **was never reachable under `--features kafka`**:
19//! `src/client/mod.rs` gated it on `#[cfg(all(feature = "kafka", feature = "rdkafka"))]`, and
20//! only `all-protocols` turned the implicit `rdkafka` feature on — so a default build linked
21//! a C library annotated in `Cargo.toml` as crashing in malloc, and every targeted
22//! `--features kafka` build silently had no Kafka client at all. Its four E2E tests were
23//! `#[ignore]`d for exactly that reason. `rdkafka` is now gone from the dependency list.
24//!
25//! # Supported surface
26//!
27//! Exactly the five APIs NetGet's broker implements, negotiated rather than assumed:
28//!
29//! | API | key | this client asks for at most | why that ceiling |
30//! |---|---|---|---|
31//! | ApiVersions  | 18 | v3 | falls back to decoding at v0 if the broker refuses, per Kafka's own rule |
32//! | Metadata     | 3  | v12 | the newest `kafka-protocol` 0.14 knows; v9+ is flexible, v10+ reports topic ids |
33//! | Produce      | 0  | v11 | the newest `kafka-protocol` 0.14 knows; v9+ is flexible |
34//! | Fetch        | 1  | v16 | v12+ is flexible; v13+ names the topic by the id Metadata reported, and a broker that reports none is fetched by name at v12 |
35//! | OffsetCommit | 8  | v9 | v8+ is flexible; a commit outside any group (generation/epoch -1) means the same at every version |
36//!
37//! The actual version used for each is `min(our ceiling, the broker's max)`, refused
38//! outright if that falls below the broker's minimum.
39//!
40//! **No consumer groups.** `FindCoordinator`, `JoinGroup`, `SyncGroup` and `Heartbeat` are
41//! implemented by neither half, so partitions are assigned manually and fetch offsets are
42//! explicit. **No `ListOffsets`**, so there is no earliest/latest resolution — a consumer
43//! starts at `start_offset` (default 0).
44//!
45//! # Why one mutex over the whole connection
46//!
47//! Kafka multiplexes requests on one TCP connection and matches replies by correlation id.
48//! Rather than run a demultiplexer task, this client holds a `tokio::sync::Mutex` for the
49//! duration of one request/response exchange, which makes the correlation id trivially
50//! unambiguous. The lock is *not* held across an LLM call — events are built, the guard is
51//! dropped, and only then is the model consulted. Every socket operation carries
52//! [`IO_TIMEOUT`], so a silent broker cannot hold the lock forever.
53
54pub mod actions;
55
56pub use actions::KafkaClientProtocol;
57
58use anyhow::{anyhow, bail, Context, Result};
59use bytes::Bytes;
60use std::collections::{HashMap, VecDeque};
61use std::net::SocketAddr;
62use std::sync::Arc;
63use std::time::Duration;
64use tokio::io::{AsyncReadExt, AsyncWriteExt};
65use tokio::net::TcpStream;
66use tokio::sync::{mpsc, Mutex};
67use tracing::{debug, error, info, trace, warn};
68
69use crate::client::kafka::actions::{
70    kafka_error_name, KAFKA_CLIENT_CONNECTED_EVENT, KAFKA_CLIENT_MESSAGE_DELIVERED_EVENT,
71    KAFKA_CLIENT_METADATA_RECEIVED_EVENT, KAFKA_CLIENT_RECORDS_RECEIVED_EVENT,
72};
73use crate::client::llm_budget::call_llm_for_client;
74use crate::llm::actions::client_trait::{Client, ClientActionResult};
75use crate::llm::ollama_client::OllamaClient;
76use crate::llm::ClientLlmResult;
77use crate::protocol::{Event, StartupParams};
78use crate::state::app_state::AppState;
79use crate::state::{ClientId, ClientStatus};
80
81// The broker's copy of the wire library and its two record-field encoding helpers. Sharing
82// them is what keeps "hex means hex" true on both sides of the same connection.
83use crate::server::kafka::kafka_protocol;
84use crate::server::kafka::{decode_field, encode_field};
85
86use kafka_protocol::messages::{
87    ApiKey, ApiVersionsRequest, ApiVersionsResponse, FetchRequest, FetchResponse, MetadataRequest,
88    MetadataResponse, OffsetCommitRequest, OffsetCommitResponse, ProduceRequest, ProduceResponse,
89    RequestHeader, ResponseHeader,
90};
91use kafka_protocol::protocol::{Decodable, Encodable, StrBytes};
92use kafka_protocol::records::{
93    Compression, Record, RecordBatchDecoder, RecordBatchEncoder, RecordEncodeOptions, TimestampType,
94};
95
96/// Every socket operation is bounded. Without this a broker that accepts a request and never
97/// answers holds the exchange mutex for the life of the process.
98const IO_TIMEOUT: Duration = Duration::from_secs(30);
99
100/// Largest response frame accepted. Mirrors the broker's own request cap; the size prefix
101/// comes off the wire and is otherwise an allocation primitive.
102const MAX_RESPONSE_BYTES: i64 = 100 * 1024 * 1024;
103
104const CLIENT_MAX_API_VERSIONS: i16 = 3;
105const CLIENT_MAX_METADATA: i16 = 12;
106const CLIENT_MAX_PRODUCE: i16 = 11;
107const CLIENT_MAX_FETCH: i16 = 16;
108const CLIENT_MAX_OFFSET_COMMIT: i16 = 9;
109/// The first Fetch version that names its topic by id rather than by name.
110const FETCH_BY_TOPIC_ID: i16 = 13;
111/// The first Metadata version that reports topic ids.
112const METADATA_WITH_TOPIC_IDS: i16 = 10;
113/// Topic ids remembered per connection.
114const MAX_TOPIC_IDS: usize = 1024;
115
116/// How many records of a batch are described to the model, and how much of each value.
117/// Matches the broker's `MAX_EVENT_RECORDS` / `MAX_EVENT_VALUE_BYTES`.
118const MAX_EVENT_RECORDS: usize = 20;
119const MAX_EVENT_VALUE_BYTES: usize = 1024;
120
121const DEFAULT_PARTITION_MAX_BYTES: i32 = 1024 * 1024;
122const DEFAULT_POLL_INTERVAL_MS: u64 = 1000;
123const MIN_POLL_INTERVAL_MS: u64 = 50;
124const DEFAULT_CLIENT_ID: &str = "netget-kafka-client";
125const DEFAULT_GROUP_ID: &str = "netget-consumer-group";
126
127/// Compression type parameters. `None` selects `kafka-protocol`'s built-in codecs (gzip,
128/// snappy, lz4, zstd are all on by default); the function types still have to be nameable.
129type Compressor = fn(&mut bytes::BytesMut, &mut Vec<u8>, Compression) -> Result<()>;
130type Decompressor = fn(&mut Bytes, Compression) -> Result<std::io::Cursor<Bytes>>;
131
132/// One partition this client is following, and where it will read from next.
133#[derive(Clone, Debug)]
134struct Assignment {
135    topic: String,
136    partition: i32,
137    next_offset: i64,
138}
139
140/// The broker connection plus everything needed to frame a request on it.
141struct KafkaConn {
142    stream: TcpStream,
143    next_correlation_id: i32,
144    /// `api_key -> (min_version, max_version)` as the broker advertised it.
145    versions: HashMap<i16, (i16, i16)>,
146    client_id: StrBytes,
147    /// Size of the last frame this connection actually put on the wire (4-byte length
148    /// prefix included). Read while the connection mutex is still held, so an injected
149    /// action can report a byte count that is genuinely its own rather than a snapshot
150    /// racing the poll loop.
151    last_request_bytes: usize,
152    /// Topic ids the broker reported, for Fetch v13+. Looked up by name on first use.
153    topic_ids: HashMap<String, uuid::Uuid>,
154}
155
156impl KafkaConn {
157    /// The broker's id for `topic`, asked for with a Metadata request the first time.
158    /// `None` when the broker reports none (Metadata below v10, or a nil id): the caller
159    /// then fetches by name, at most v12.
160    async fn topic_id(&mut self, topic: &str) -> Result<Option<uuid::Uuid>> {
161        if let Some(id) = self.topic_ids.get(topic) {
162            return Ok(Some(*id));
163        }
164        use kafka_protocol::messages::metadata_request::MetadataRequestTopic;
165        let version = self.version_for(ApiKey::Metadata, CLIENT_MAX_METADATA)?;
166        if version < METADATA_WITH_TOPIC_IDS {
167            return Ok(None);
168        }
169        let response: MetadataResponse = self
170            .exchange(
171                ApiKey::Metadata,
172                version,
173                &MetadataRequest::default()
174                    .with_topics(Some(vec![MetadataRequestTopic::default()
175                        .with_name(Some(StrBytes::from_string(topic.to_string()).into()))])),
176            )
177            .await
178            .context("Metadata request for a topic id failed")?;
179        self.remember_topic_ids(&response);
180        Ok(self.topic_ids.get(topic).copied())
181    }
182
183    /// Keep the non-nil ids of the topics a Metadata response described without error.
184    fn remember_topic_ids(&mut self, response: &MetadataResponse) {
185        for t in &response.topics {
186            let Some(name) = t.name.as_ref() else {
187                continue;
188            };
189            if t.error_code != 0 || t.topic_id.is_nil() {
190                continue;
191            }
192            if self.topic_ids.len() >= MAX_TOPIC_IDS && !self.topic_ids.contains_key(name.as_str())
193            {
194                self.topic_ids.clear();
195            }
196            self.topic_ids.insert(name.to_string(), t.topic_id);
197        }
198    }
199
200    /// The version to use for `api_key`: `min(our ceiling, the broker's max)`, refused if
201    /// that lands below the broker's minimum or the broker does not implement the API at all.
202    ///
203    /// Guessing a version here is not a small mistake: a body encoded at a version the broker
204    /// does not implement is not rejected cleanly, it is misparsed.
205    fn version_for(&self, api_key: ApiKey, our_max: i16) -> Result<i16> {
206        let (broker_min, broker_max) =
207            *self.versions.get(&(api_key as i16)).with_context(|| {
208                format!(
209                    "broker does not advertise {api_key:?} (key {})",
210                    api_key as i16
211                )
212            })?;
213        let chosen = our_max.min(broker_max);
214        if chosen < broker_min {
215            bail!(
216                "broker supports {api_key:?} v{broker_min}-v{broker_max}, but this client \
217                 implements at most v{our_max}"
218            );
219        }
220        Ok(chosen)
221    }
222
223    fn build_request<B: Encodable>(
224        &mut self,
225        api_key: ApiKey,
226        api_version: i16,
227        body: &B,
228    ) -> Result<(i32, Vec<u8>)> {
229        let correlation_id = self.next_correlation_id;
230        self.next_correlation_id = self.next_correlation_id.wrapping_add(1);
231
232        let mut buf = Vec::new();
233        RequestHeader::default()
234            .with_request_api_key(api_key as i16)
235            .with_request_api_version(api_version)
236            .with_correlation_id(correlation_id)
237            .with_client_id(Some(self.client_id.clone()))
238            .encode(&mut buf, api_key.request_header_version(api_version))
239            .with_context(|| format!("encoding {api_key:?} v{api_version} request header"))?;
240        body.encode(&mut buf, api_version)
241            .with_context(|| format!("encoding {api_key:?} v{api_version} request body"))?;
242        Ok((correlation_id, buf))
243    }
244
245    async fn write_frame(&mut self, body: &[u8]) -> Result<()> {
246        let size = i32::try_from(body.len()).context("request does not fit in an i32 frame")?;
247        tokio::time::timeout(IO_TIMEOUT, self.stream.write_all(&size.to_be_bytes()))
248            .await
249            .context("timed out writing Kafka frame size")??;
250        tokio::time::timeout(IO_TIMEOUT, self.stream.write_all(body))
251            .await
252            .context("timed out writing Kafka frame body")??;
253        self.last_request_bytes = 4 + body.len();
254        Ok(())
255    }
256
257    async fn read_frame(&mut self) -> Result<Vec<u8>> {
258        let mut size = [0u8; 4];
259        tokio::time::timeout(IO_TIMEOUT, self.stream.read_exact(&mut size))
260            .await
261            .context("timed out reading Kafka response size")??;
262        // Validated in i64 so neither sign extension nor the length itself can wrap.
263        let announced = i64::from(i32::from_be_bytes(size));
264        if announced <= 0 || announced > MAX_RESPONSE_BYTES {
265            bail!("broker announced an implausible response size of {announced} bytes");
266        }
267        let mut buf = vec![0u8; announced as usize];
268        tokio::time::timeout(IO_TIMEOUT, self.stream.read_exact(&mut buf))
269            .await
270            .context("timed out reading Kafka response body")??;
271        Ok(buf)
272    }
273
274    /// One request, one response, matched on correlation id.
275    async fn exchange<B: Encodable, R: Decodable>(
276        &mut self,
277        api_key: ApiKey,
278        api_version: i16,
279        body: &B,
280    ) -> Result<R> {
281        let (correlation_id, request) = self.build_request(api_key, api_version, body)?;
282        self.write_frame(&request).await?;
283        let response = self.read_frame().await?;
284
285        let mut cursor = std::io::Cursor::new(response.as_slice());
286        let header =
287            ResponseHeader::decode(&mut cursor, api_key.response_header_version(api_version))
288                .with_context(|| format!("decoding {api_key:?} v{api_version} response header"))?;
289        if header.correlation_id != correlation_id {
290            bail!(
291                "broker answered {api_key:?} with correlation id {} instead of {correlation_id}; \
292                 the connection is desynchronised",
293                header.correlation_id
294            );
295        }
296        R::decode(&mut cursor, api_version)
297            .with_context(|| format!("decoding {api_key:?} v{api_version} response body"))
298    }
299
300    /// Send a request and deliberately read nothing back. Only correct for `acks=0` Produce,
301    /// where the broker is specified to write no reply at all.
302    async fn send_only<B: Encodable>(
303        &mut self,
304        api_key: ApiKey,
305        api_version: i16,
306        body: &B,
307    ) -> Result<()> {
308        let (_, request) = self.build_request(api_key, api_version, body)?;
309        self.write_frame(&request).await
310    }
311}
312
313/// Everything one connected client needs to run its event loop.
314#[derive(Clone)]
315struct Session {
316    conn: Arc<Mutex<KafkaConn>>,
317    assignments: Arc<Mutex<Vec<Assignment>>>,
318    memory: Arc<Mutex<String>>,
319    client_id: ClientId,
320    app_state: Arc<AppState>,
321    llm_client: OllamaClient,
322    status_tx: mpsc::UnboundedSender<String>,
323    default_group_id: String,
324    remote_addr: String,
325}
326
327/// Kafka client that connects to a broker.
328pub struct KafkaClient;
329
330impl KafkaClient {
331    /// Connect to a broker, negotiate, fetch metadata, then hand the session to the model.
332    pub async fn connect_with_llm_actions(
333        remote_addr: String,
334        llm_client: OllamaClient,
335        app_state: Arc<AppState>,
336        status_tx: mpsc::UnboundedSender<String>,
337        client_id: ClientId,
338        startup_params: Option<StartupParams>,
339    ) -> Result<SocketAddr> {
340        // Parameters are validated before the socket is opened, so a bad value fails the
341        // connect with a message naming the key instead of half-starting a client.
342        let params = startup_params.as_ref();
343        let client_id_str = params
344            .map(|p| p.get_optional_string("client_id"))
345            .transpose()?
346            .flatten()
347            .unwrap_or_else(|| DEFAULT_CLIENT_ID.to_string());
348
349        let topics: Vec<String> = params
350            .map(|p| p.get_optional_array("topics"))
351            .transpose()?
352            .flatten()
353            .map(|arr| {
354                arr.iter()
355                    .filter_map(|v| v.as_str())
356                    .filter(|s| !s.is_empty())
357                    .map(|s| s.to_string())
358                    .collect()
359            })
360            .unwrap_or_default();
361
362        let partition_raw = params
363            .map(|p| p.get_optional_i64("partition"))
364            .transpose()?
365            .flatten()
366            .unwrap_or(0);
367        if !(0..=actions::MAX_PARTITION).contains(&partition_raw) {
368            bail!(
369                "Kafka client 'partition' must be between 0 and {}, got {partition_raw}",
370                actions::MAX_PARTITION
371            );
372        }
373        let partition = partition_raw as i32;
374
375        let start_offset = params
376            .map(|p| p.get_optional_i64("start_offset"))
377            .transpose()?
378            .flatten()
379            .unwrap_or(0);
380        if start_offset < 0 {
381            bail!("Kafka client 'start_offset' must not be negative, got {start_offset}");
382        }
383
384        let poll_interval_raw = params
385            .map(|p| p.get_optional_u64("poll_interval_ms"))
386            .transpose()?
387            .flatten()
388            .unwrap_or(DEFAULT_POLL_INTERVAL_MS);
389        if poll_interval_raw < MIN_POLL_INTERVAL_MS {
390            bail!(
391                "Kafka client 'poll_interval_ms' must be at least {MIN_POLL_INTERVAL_MS}, got \
392                 {poll_interval_raw}"
393            );
394        }
395        let poll_interval = Duration::from_millis(poll_interval_raw);
396
397        let group_id = params
398            .map(|p| p.get_optional_string("group_id"))
399            .transpose()?
400            .flatten()
401            .unwrap_or_else(|| DEFAULT_GROUP_ID.to_string());
402
403        info!(
404            "Kafka client {} connecting to {} (client_id={}, topics={:?})",
405            client_id, remote_addr, client_id_str, topics
406        );
407
408        let stream = tokio::time::timeout(IO_TIMEOUT, TcpStream::connect(&remote_addr))
409            .await
410            .with_context(|| format!("timed out connecting to Kafka broker {remote_addr}"))?
411            .with_context(|| format!("failed to connect to Kafka broker {remote_addr}"))?;
412        let local_addr = stream.local_addr()?;
413
414        let mut conn = KafkaConn {
415            stream,
416            next_correlation_id: 1,
417            versions: HashMap::new(),
418            client_id: StrBytes::from_string(client_id_str.clone()),
419            last_request_bytes: 0,
420            topic_ids: HashMap::new(),
421        };
422
423        // ---- ApiVersions ------------------------------------------------------------
424        conn.versions = Self::negotiate_api_versions(&mut conn).await?;
425
426        let negotiated = serde_json::json!({
427            "metadata": conn.version_for(ApiKey::Metadata, CLIENT_MAX_METADATA)?,
428            "produce": conn.version_for(ApiKey::Produce, CLIENT_MAX_PRODUCE)?,
429            "fetch": conn.version_for(ApiKey::Fetch, CLIENT_MAX_FETCH)?,
430            "offset_commit": conn.version_for(ApiKey::OffsetCommit, CLIENT_MAX_OFFSET_COMMIT)?,
431        });
432
433        // ---- Metadata ---------------------------------------------------------------
434        let requested: Option<Vec<String>> = if topics.is_empty() {
435            None
436        } else {
437            Some(topics.clone())
438        };
439        let mut metadata = Self::request_metadata(&mut conn, requested.as_deref()).await?;
440
441        info!(
442            "Kafka client {} negotiated {} and read metadata for {} topic(s)",
443            client_id,
444            negotiated,
445            metadata
446                .get("topics")
447                .and_then(|t| t.as_array())
448                .map(|a| a.len())
449                .unwrap_or(0)
450        );
451
452        app_state
453            .update_client_status(client_id, ClientStatus::Connected)
454            .await;
455        let _ = status_tx.send(format!(
456            "[CLIENT] Kafka client {client_id} connected to {remote_addr}"
457        ));
458        let _ = status_tx.send("__UPDATE_UI__".to_string());
459
460        app_state
461            .with_client_mut(client_id, |client_inst| {
462                client_inst
463                    .set_protocol_field("brokers".to_string(), serde_json::json!(remote_addr));
464                client_inst.set_protocol_field(
465                    "kafka_client_id".to_string(),
466                    serde_json::json!(client_id_str),
467                );
468                client_inst.set_protocol_field("api_versions".to_string(), negotiated.clone());
469            })
470            .await;
471
472        let assignments: Vec<Assignment> = topics
473            .iter()
474            .map(|topic| Assignment {
475                topic: topic.clone(),
476                partition,
477                next_offset: start_offset,
478            })
479            .collect();
480        let polling = !assignments.is_empty();
481
482        let memory = app_state
483            .get_memory_for_client(client_id)
484            .await
485            .unwrap_or_default();
486
487        let session = Session {
488            conn: Arc::new(Mutex::new(conn)),
489            assignments: Arc::new(Mutex::new(assignments)),
490            memory: Arc::new(Mutex::new(memory)),
491            client_id,
492            app_state: app_state.clone(),
493            llm_client,
494            status_tx: status_tx.clone(),
495            default_group_id: group_id,
496            remote_addr: remote_addr.clone(),
497        };
498
499        // The connected event and everything after it run in a task, so `connect` returns as
500        // soon as the transport is genuinely up rather than blocking startup on the model.
501        let connected_data = {
502            let obj = metadata
503                .as_object_mut()
504                .expect("request_metadata builds an object");
505            obj.insert("remote_addr".to_string(), serde_json::json!(remote_addr));
506            obj.insert("api_versions".to_string(), negotiated);
507            metadata
508        };
509
510        // The dashboard's `[ send ]` channel, registered BEFORE the connected-event LLM call
511        // below. A dashboard-created client defaults to a `*` -> manual rule, so that call can
512        // park for minutes waiting for a human; registering after it would leave the rail
513        // reading "no command channel" for the whole park.
514        let command_rx =
515            crate::client::command_support::register_command_channel(&app_state, client_id).await;
516        let command_session = session.clone();
517        let command_task = tokio::spawn(async move {
518            command_session.command_loop(command_rx).await;
519        });
520        app_state
521            .register_client_task(client_id, command_task)
522            .await;
523
524        let task_registrar = app_state.clone();
525        let task_handle = tokio::spawn(async move {
526            let stop = session
527                .drive(Event::new(&KAFKA_CLIENT_CONNECTED_EVENT, connected_data))
528                .await;
529
530            if stop {
531                session.close("handler asked to disconnect").await;
532                return;
533            }
534            if !polling {
535                debug!(
536                    "Kafka client {} has no topics to poll; the connection stays open for \
537                     scheduled tasks and further instructions",
538                    session.client_id
539                );
540                return;
541            }
542            session.poll_loop(poll_interval).await;
543        });
544        task_registrar
545            .register_client_task(client_id, task_handle)
546            .await;
547
548        Ok(local_addr)
549    }
550
551    /// Ask the broker what it speaks.
552    ///
553    /// Kafka's own negotiation rule: a broker that does not implement the requested
554    /// ApiVersions version answers `UNSUPPORTED_VERSION` *plus the supported-API table*,
555    /// encoded at v0 so the client can always read it. Both cases are handled from the same
556    /// bytes — decoding at v3 first, then at v0 — because the reply is the only thing that
557    /// says which version to step down to.
558    async fn negotiate_api_versions(conn: &mut KafkaConn) -> Result<HashMap<i16, (i16, i16)>> {
559        let request = ApiVersionsRequest::default()
560            .with_client_software_name(StrBytes::from_static_str("netget"))
561            .with_client_software_version(StrBytes::from_static_str(env!("CARGO_PKG_VERSION")));
562
563        let (correlation_id, bytes) =
564            conn.build_request(ApiKey::ApiVersions, CLIENT_MAX_API_VERSIONS, &request)?;
565        conn.write_frame(&bytes).await?;
566        let response = conn.read_frame().await?;
567
568        let decoded = Self::decode_api_versions(&response, CLIENT_MAX_API_VERSIONS, correlation_id)
569            .and_then(|(code, body)| {
570                if code == 0 {
571                    Ok(body)
572                } else {
573                    Err(anyhow!(
574                        "ApiVersions v{CLIENT_MAX_API_VERSIONS} refused with error {code}"
575                    ))
576                }
577            });
578
579        let body = match decoded {
580            Ok(body) => body,
581            Err(first) => {
582                debug!("Kafka client stepping ApiVersions down to v0: {}", first);
583                let (code, body) = Self::decode_api_versions(&response, 0, correlation_id)
584                    .context("broker's ApiVersions reply is unreadable at v3 and at v0")?;
585                if code != 0 && body.api_keys.is_empty() {
586                    bail!(
587                        "broker refused ApiVersions with error {code} ({}) and sent no supported-API \
588                         table, so there is no version to step down to",
589                        kafka_error_name(code)
590                    );
591                }
592                body
593            }
594        };
595
596        if body.api_keys.is_empty() {
597            bail!("broker advertised no APIs at all");
598        }
599
600        Ok(body
601            .api_keys
602            .iter()
603            .map(|a| (a.api_key, (a.min_version, a.max_version)))
604            .collect())
605    }
606
607    fn decode_api_versions(
608        bytes: &[u8],
609        api_version: i16,
610        expected_correlation_id: i32,
611    ) -> Result<(i16, ApiVersionsResponse)> {
612        let mut cursor = std::io::Cursor::new(bytes);
613        let header = ResponseHeader::decode(
614            &mut cursor,
615            ApiKey::ApiVersions.response_header_version(api_version),
616        )
617        .context("decoding ApiVersions response header")?;
618        if header.correlation_id != expected_correlation_id {
619            bail!(
620                "broker answered ApiVersions with correlation id {} instead of {}",
621                header.correlation_id,
622                expected_correlation_id
623            );
624        }
625        let body = ApiVersionsResponse::decode(&mut cursor, api_version)
626            .context("decoding ApiVersions response body")?;
627        Ok((body.error_code, body))
628    }
629
630    /// Metadata request → the structured shape a handler sees.
631    async fn request_metadata(
632        conn: &mut KafkaConn,
633        topics: Option<&[String]>,
634    ) -> Result<serde_json::Value> {
635        use kafka_protocol::messages::metadata_request::MetadataRequestTopic;
636
637        let version = conn.version_for(ApiKey::Metadata, CLIENT_MAX_METADATA)?;
638
639        let named: Option<Vec<MetadataRequestTopic>> = match topics {
640            Some(names) => Some(
641                names
642                    .iter()
643                    .map(|n| {
644                        MetadataRequestTopic::default()
645                            .with_name(Some(StrBytes::from_string(n.clone()).into()))
646                    })
647                    .collect(),
648            ),
649            // v0 has no way to express "null topics"; there, an empty list means "everything".
650            None if version == 0 => Some(Vec::new()),
651            None => None,
652        };
653
654        let response: MetadataResponse = conn
655            .exchange(
656                ApiKey::Metadata,
657                version,
658                &MetadataRequest::default().with_topics(named),
659            )
660            .await
661            .context("Metadata request failed")?;
662        // Every Metadata answer from v10 carries topic ids; remember them, so a Fetch v13+
663        // on a topic already described costs no extra round trip.
664        conn.remember_topic_ids(&response);
665
666        Ok(metadata_to_json(&response))
667    }
668}
669
670impl Session {
671    /// Run a handler for `event`, execute what it returns, and keep going for any event those
672    /// actions produced. Returns true when the session should be torn down.
673    ///
674    /// Iterative, not recursive: the DNS client reached 211 model calls and then overflowed
675    /// the stack doing this with recursion (`IMPROVEMENTS.md` item 49). The LLM budget in
676    /// [`crate::client::llm_budget`] bounds the total; this shape bounds the stack.
677    async fn drive(&self, first: Event) -> bool {
678        let Some(instruction) = self
679            .app_state
680            .get_instruction_for_client(self.client_id)
681            .await
682        else {
683            return false;
684        };
685
686        let mut queue: VecDeque<Event> = VecDeque::from([first]);
687        while let Some(event) = queue.pop_front() {
688            let memory = self.memory.lock().await.clone();
689            let protocol = KafkaClientProtocol::new();
690
691            let result = call_llm_for_client(
692                &self.llm_client,
693                &self.app_state,
694                self.client_id.to_string(),
695                &instruction,
696                &memory,
697                Some(&event),
698                &protocol,
699                &self.status_tx,
700            )
701            .await;
702
703            let ClientLlmResult {
704                actions,
705                memory_updates,
706            } = match result {
707                Ok(r) => r,
708                Err(e) => {
709                    error!(
710                        "Kafka client {} could not handle {}: {}",
711                        self.client_id,
712                        event.id(),
713                        e
714                    );
715                    let _ = self
716                        .status_tx
717                        .send(format!("[CLIENT] Kafka handler failed: {e}"));
718                    return false;
719                }
720            };
721            if let Some(mem) = memory_updates {
722                *self.memory.lock().await = mem;
723            }
724
725            for action in actions {
726                match protocol.execute_action(action) {
727                    Ok(ClientActionResult::Disconnect) => return true,
728                    Ok(ClientActionResult::WaitForMore) | Ok(ClientActionResult::NoAction) => {}
729                    Ok(ClientActionResult::Custom { name, data }) => {
730                        match self.perform(&name, &data).await {
731                            Ok((Some(follow_up), _)) => queue.push_back(follow_up),
732                            Ok((None, _)) => {}
733                            Err(e) => {
734                                // A failed exchange is reported, never papered over with a
735                                // plausible-looking success the handler would act on.
736                                error!(
737                                    "Kafka client {} action '{}' failed: {}",
738                                    self.client_id, name, e
739                                );
740                                let _ = self
741                                    .status_tx
742                                    .send(format!("[CLIENT] Kafka {name} failed: {e}"));
743                                if is_fatal(&e) {
744                                    self.mark_error(&e.to_string()).await;
745                                    return true;
746                                }
747                            }
748                        }
749                    }
750                    Ok(other) => {
751                        warn!(
752                            "Kafka client {} ignoring unsupported action result {:?}",
753                            self.client_id, other
754                        );
755                    }
756                    Err(e) => {
757                        error!("Kafka client {} rejected an action: {}", self.client_id, e);
758                        let _ = self
759                            .status_tx
760                            .send(format!("[CLIENT] Kafka action rejected: {e}"));
761                    }
762                }
763            }
764        }
765        false
766    }
767
768    /// Drain injected commands (the dashboard's `[ send ]`) until the channel closes - which
769    /// happens when the client is removed or [`Self::close`] drops the handle - or an injected
770    /// `disconnect` ends the session.
771    ///
772    /// Kafka frames are read with `read_exact`, which is **not** cancellation-safe, so this is
773    /// its own task rather than a `select!` arm in the poll loop. The two are serialised by
774    /// the connection mutex the exchange already takes, which is also what makes the byte
775    /// count reported below unambiguously this command's own request.
776    ///
777    /// Injected actions go through [`Self::perform`] - the exact function the LLM path uses -
778    /// so there is no second copy of the wire encoding. A follow-up event (a Produce ack, a
779    /// Fetch result, Metadata) is handed to [`Self::drive`] in its own task rather than
780    /// inline, so a handler that parks for a human cannot block the next injected command.
781    async fn command_loop(
782        &self,
783        mut command_rx: mpsc::Receiver<crate::state::client_handles::ClientCommand>,
784    ) {
785        use crate::llm::actions::protocol_trait::Protocol;
786        use crate::state::client_handles::ClientSendOutcome;
787        use crate::state::AccessLogOwner;
788
789        let protocol = KafkaClientProtocol::new();
790
791        while let Some(command) = command_rx.recv().await {
792            let action = command.action.clone();
793            let mut follow_up: Option<Event> = None;
794            let mut disconnect = false;
795
796            let outcome: Result<ClientSendOutcome> = match protocol.execute_action(action.clone()) {
797                Err(e) => Ok(ClientSendOutcome::Rejected {
798                    error: e.to_string(),
799                }),
800                Ok(ClientActionResult::Disconnect) => {
801                    disconnect = true;
802                    Ok(ClientSendOutcome::Disconnected)
803                }
804                Ok(ClientActionResult::WaitForMore) => Ok(ClientSendOutcome::Executed {
805                    detail: "wait_for_more: nothing was asked of the broker".to_string(),
806                }),
807                Ok(ClientActionResult::NoAction) => Ok(ClientSendOutcome::Executed {
808                    detail: "no_action".to_string(),
809                }),
810                Ok(ClientActionResult::Custom { name, data }) => {
811                    match self.perform(&name, &data).await {
812                        Ok((event, wrote)) => {
813                            follow_up = event;
814                            Ok(ClientSendOutcome::Sent { bytes_sent: wrote })
815                        }
816                        Err(e) => {
817                            // A refused exchange is reported as a failure, never as a
818                            // plausible-looking `Sent` the operator would trust.
819                            if is_fatal(&e) {
820                                self.mark_error(&e.to_string()).await;
821                            }
822                            Err(e.context(format!("injected Kafka action '{name}'")))
823                        }
824                    }
825                }
826                Ok(other) => Ok(ClientSendOutcome::Executed {
827                    detail: format!("unsupported action result {other:?}"),
828                }),
829            };
830
831            let outcome_json = match &outcome {
832                Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
833                Err(e) => serde_json::json!({"error": e.to_string()}),
834            };
835            self.app_state
836                .record_access_log(
837                    AccessLogOwner::Client(self.client_id.as_u32()),
838                    protocol.protocol_name(),
839                    None,
840                    "injected_action",
841                    action,
842                    vec![outcome_json],
843                )
844                .await;
845
846            if let Err(e) = &outcome {
847                error!(
848                    "Kafka client {} injected action failed: {}",
849                    self.client_id, e
850                );
851                let _ = self
852                    .status_tx
853                    .send(format!("[CLIENT] Kafka injected action failed: {e}"));
854            }
855            let _ = self.status_tx.send("__UPDATE_UI__".to_string());
856            crate::client::command_support::reply(command, outcome);
857
858            if disconnect {
859                self.close("injected disconnect").await;
860                break;
861            }
862
863            if let Some(event) = follow_up {
864                let session = self.clone();
865                let handle = tokio::spawn(async move {
866                    if session.drive(event).await {
867                        session.close("handler asked to disconnect").await;
868                    }
869                });
870                self.app_state
871                    .register_client_task(self.client_id, handle)
872                    .await;
873            }
874        }
875
876        // The channel only ends here, so this is the one exit: drop the handle so a late
877        // send fails fast instead of queueing into a dead session.
878        self.app_state.remove_client_handle(self.client_id).await;
879    }
880
881    /// Perform one action against the broker. Returns the event it produced, if any, and
882    /// the number of request bytes this action put on the wire.
883    ///
884    /// The byte count is read out while the connection mutex is still held by the operation
885    /// that wrote it, so an injected command can report `ClientSendOutcome::Sent` with a
886    /// figure that is genuinely its own request rather than a counter shared with the poll
887    /// loop.
888    async fn perform(
889        &self,
890        name: &str,
891        data: &serde_json::Value,
892    ) -> Result<(Option<Event>, usize)> {
893        match name {
894            "kafka_produce" => self.produce(data).await,
895            "kafka_fetch" => {
896                let topic = str_field(data, "topic")?;
897                let partition = i64_field(data, "partition")? as i32;
898                let offset = match data.get("offset").and_then(|v| v.as_i64()) {
899                    Some(o) => o,
900                    None => self.tracked_offset(&topic, partition).await.unwrap_or(0),
901                };
902                let max_bytes = data
903                    .get("max_bytes")
904                    .and_then(|v| v.as_i64())
905                    .unwrap_or(i64::from(DEFAULT_PARTITION_MAX_BYTES))
906                    as i32;
907                self.fetch(&topic, partition, offset, max_bytes).await
908            }
909            "kafka_metadata" => {
910                let topics: Option<Vec<String>> =
911                    data.get("topics").and_then(|v| v.as_array()).map(|a| {
912                        a.iter()
913                            .filter_map(|v| v.as_str())
914                            .map(|s| s.to_string())
915                            .collect()
916                    });
917                let mut conn = self.conn.lock().await;
918                let metadata = KafkaClient::request_metadata(&mut conn, topics.as_deref()).await?;
919                let wrote = conn.last_request_bytes;
920                drop(conn);
921                Ok((
922                    Some(Event::new(&KAFKA_CLIENT_METADATA_RECEIVED_EVENT, metadata)),
923                    wrote,
924                ))
925            }
926            "kafka_commit" => self.commit(data).await.map(|wrote| (None, wrote)),
927            other => Err(anyhow!("unknown Kafka client action result '{other}'")),
928        }
929    }
930
931    /// Publish one record.
932    async fn produce(&self, data: &serde_json::Value) -> Result<(Option<Event>, usize)> {
933        use kafka_protocol::messages::produce_request::{PartitionProduceData, TopicProduceData};
934
935        let topic = str_field(data, "topic")?;
936        let partition = i64_field(data, "partition")? as i32;
937        let acks = i64_field(data, "acks")? as i16;
938
939        // The declared encoding is honoured here, by the same decoder the broker uses for its
940        // own outbound records. An action documented as accepting hex that puts literal ASCII
941        // on the wire is the defect this codebase names as its reference case.
942        let key = decode_field(
943            data.get("key"),
944            data.get("key_encoding")
945                .and_then(|v| v.as_str())
946                .unwrap_or("utf8"),
947        )
948        .context("record key")?;
949        let value = decode_field(
950            data.get("value"),
951            data.get("value_encoding")
952                .and_then(|v| v.as_str())
953                .unwrap_or("utf8"),
954        )
955        .context("record value")?;
956
957        let batch = encode_single_record(key, value)?;
958
959        let request = ProduceRequest::default()
960            .with_acks(acks)
961            .with_timeout_ms(IO_TIMEOUT.as_millis() as i32)
962            .with_topic_data(vec![TopicProduceData::default()
963                .with_name(StrBytes::from_string(topic.clone()).into())
964                .with_partition_data(vec![PartitionProduceData::default()
965                    .with_index(partition)
966                    .with_records(Some(batch))])]);
967
968        let mut conn = self.conn.lock().await;
969        let version = conn.version_for(ApiKey::Produce, CLIENT_MAX_PRODUCE)?;
970
971        if acks == 0 {
972            // Kafka specifies no reply for acks=0. Waiting for one would hang the exchange
973            // until IO_TIMEOUT and then look like a broker failure.
974            conn.send_only(ApiKey::Produce, version, &request).await?;
975            let wrote = conn.last_request_bytes;
976            drop(conn);
977            info!(
978                "Kafka client {} produced to {}/{} with acks=0 (no acknowledgement requested)",
979                self.client_id, topic, partition
980            );
981            let _ = self.status_tx.send(format!(
982                "[CLIENT] Kafka produced to {topic}/{partition} (acks=0, unacknowledged)"
983            ));
984            return Ok((None, wrote));
985        }
986
987        let response: ProduceResponse = conn.exchange(ApiKey::Produce, version, &request).await?;
988        let wrote = conn.last_request_bytes;
989        drop(conn);
990
991        let topic_response = response
992            .responses
993            .iter()
994            .find(|t| t.name.as_str() == topic)
995            .or_else(|| response.responses.first())
996            .context("Produce response named no topic")?;
997        let partition_response = topic_response
998            .partition_responses
999            .iter()
1000            .find(|p| p.index == partition)
1001            .or_else(|| topic_response.partition_responses.first())
1002            .context("Produce response named no partition")?;
1003
1004        let error_code = partition_response.error_code;
1005        let delivered = error_code == 0;
1006        if delivered {
1007            info!(
1008                "Kafka client {} produced to {}/{} at offset {}",
1009                self.client_id, topic, partition, partition_response.base_offset
1010            );
1011            let _ = self.status_tx.send(format!(
1012                "[CLIENT] Kafka produced to {topic}/{partition} at offset {}",
1013                partition_response.base_offset
1014            ));
1015        } else {
1016            // A rejected write is an error, not a quieter kind of success.
1017            error!(
1018                "Kafka client {} produce to {}/{} rejected: {} ({})",
1019                self.client_id,
1020                topic,
1021                partition,
1022                error_code,
1023                kafka_error_name(error_code)
1024            );
1025            let _ = self.status_tx.send(format!(
1026                "[CLIENT] Kafka produce to {topic}/{partition} rejected: {} ({error_code})",
1027                kafka_error_name(error_code)
1028            ));
1029        }
1030
1031        Ok((
1032            Some(Event::new(
1033                &KAFKA_CLIENT_MESSAGE_DELIVERED_EVENT,
1034                serde_json::json!({
1035                    "topic": topic,
1036                    "partition": partition,
1037                    "base_offset": partition_response.base_offset,
1038                    "error_code": error_code,
1039                    "error_name": kafka_error_name(error_code),
1040                    "delivered": delivered,
1041                    "api_version": version,
1042                }),
1043            )),
1044            wrote,
1045        ))
1046    }
1047
1048    /// Read one partition. Returns an event only when records actually came back, so an idle
1049    /// poll costs no model call.
1050    async fn fetch(
1051        &self,
1052        topic: &str,
1053        partition: i32,
1054        offset: i64,
1055        max_bytes: i32,
1056    ) -> Result<(Option<Event>, usize)> {
1057        use kafka_protocol::messages::fetch_request::{FetchPartition, FetchTopic};
1058
1059        let mut conn = self.conn.lock().await;
1060        let mut version = conn.version_for(ApiKey::Fetch, CLIENT_MAX_FETCH)?;
1061        // From v13 a fetch names its topic by the id the broker gave it; without one, fetch
1062        // by name at the last version that allows it.
1063        let mut fetch_topic = FetchTopic::default();
1064        if version >= FETCH_BY_TOPIC_ID {
1065            match conn.topic_id(topic).await? {
1066                Some(id) => fetch_topic = fetch_topic.with_topic_id(id),
1067                None => version = conn.version_for(ApiKey::Fetch, FETCH_BY_TOPIC_ID - 1)?,
1068            }
1069        }
1070        if version < FETCH_BY_TOPIC_ID {
1071            fetch_topic = fetch_topic.with_topic(StrBytes::from_string(topic.to_string()).into());
1072        }
1073        let request = FetchRequest::default()
1074            .with_max_wait_ms(500)
1075            .with_min_bytes(1)
1076            .with_max_bytes(max_bytes)
1077            .with_topics(vec![fetch_topic.with_partitions(vec![
1078                FetchPartition::default()
1079                    .with_partition(partition)
1080                    .with_fetch_offset(offset)
1081                    .with_partition_max_bytes(max_bytes),
1082            ])]);
1083        let response: FetchResponse = conn.exchange(ApiKey::Fetch, version, &request).await?;
1084        let wrote = conn.last_request_bytes;
1085        drop(conn);
1086
1087        if response.error_code != 0 {
1088            bail!(
1089                "Fetch failed at the response level: {} ({})",
1090                response.error_code,
1091                kafka_error_name(response.error_code)
1092            );
1093        }
1094
1095        let Some(topic_response) = response.responses.first() else {
1096            debug!(
1097                "Kafka client {} fetched {}/{}: broker returned no topic",
1098                self.client_id, topic, partition
1099            );
1100            return Ok((None, wrote));
1101        };
1102        let Some(partition_data) = topic_response.partitions.first() else {
1103            return Ok((None, wrote));
1104        };
1105
1106        if partition_data.error_code != 0 {
1107            bail!(
1108                "Fetch of {topic}/{partition} failed: {} ({})",
1109                partition_data.error_code,
1110                kafka_error_name(partition_data.error_code)
1111            );
1112        }
1113
1114        let records = match partition_data.records.as_ref() {
1115            Some(raw) if !raw.is_empty() => {
1116                let owned = Bytes::copy_from_slice(raw.as_ref());
1117                let mut cursor = std::io::Cursor::new(owned);
1118                RecordBatchDecoder::decode_with_custom_compression::<_, Decompressor>(
1119                    &mut cursor,
1120                    None,
1121                )
1122                .context("broker's record batch did not decode")?
1123            }
1124            _ => Vec::new(),
1125        };
1126
1127        if records.is_empty() {
1128            trace!(
1129                "Kafka client {} fetched {}/{} from offset {}: nothing new",
1130                self.client_id,
1131                topic,
1132                partition,
1133                offset
1134            );
1135            return Ok((None, wrote));
1136        }
1137
1138        let last_offset = records.iter().map(|r| r.offset).max().unwrap_or(offset);
1139        let next_offset = last_offset + 1;
1140        self.set_tracked_offset(topic, partition, next_offset).await;
1141
1142        info!(
1143            "Kafka client {} fetched {} record(s) from {}/{} at offset {}",
1144            self.client_id,
1145            records.len(),
1146            topic,
1147            partition,
1148            offset
1149        );
1150        let _ = self.status_tx.send(format!(
1151            "[CLIENT] Kafka received {} record(s) from {topic}/{partition}",
1152            records.len()
1153        ));
1154
1155        Ok((
1156            Some(Event::new(
1157                &KAFKA_CLIENT_RECORDS_RECEIVED_EVENT,
1158                serde_json::json!({
1159                    "topic": topic,
1160                    "partition": partition,
1161                    "api_version": version,
1162                    "high_watermark": partition_data.high_watermark,
1163                    "next_offset": next_offset,
1164                    "record_count": records.len(),
1165                    "records": records.iter().take(MAX_EVENT_RECORDS).map(record_to_json).collect::<Vec<_>>(),
1166                }),
1167            )),
1168            wrote,
1169        ))
1170    }
1171
1172    /// Commit an offset. No event: the answer is a bare error code, and raising an event for
1173    /// it would double the model round trips a consumer needs per batch.
1174    async fn commit(&self, data: &serde_json::Value) -> Result<usize> {
1175        use kafka_protocol::messages::offset_commit_request::{
1176            OffsetCommitRequestPartition, OffsetCommitRequestTopic,
1177        };
1178
1179        let topic = str_field(data, "topic")?;
1180        let partition = i64_field(data, "partition")? as i32;
1181        let offset = i64_field(data, "offset")?;
1182        let group_id = data
1183            .get("group_id")
1184            .and_then(|v| v.as_str())
1185            .unwrap_or(&self.default_group_id)
1186            .to_string();
1187
1188        let request = OffsetCommitRequest::default()
1189            .with_group_id(StrBytes::from_string(group_id.clone()).into())
1190            // -1 / empty: this client is not a group member, and claiming a generation it
1191            // never joined would be rejected by any broker that implements groups.
1192            .with_generation_id_or_member_epoch(-1)
1193            .with_topics(vec![OffsetCommitRequestTopic::default()
1194                .with_name(StrBytes::from_string(topic.clone()).into())
1195                .with_partitions(vec![OffsetCommitRequestPartition::default()
1196                    .with_partition_index(partition)
1197                    .with_committed_offset(offset)])]);
1198
1199        let mut conn = self.conn.lock().await;
1200        let version = conn.version_for(ApiKey::OffsetCommit, CLIENT_MAX_OFFSET_COMMIT)?;
1201        let response: OffsetCommitResponse = conn
1202            .exchange(ApiKey::OffsetCommit, version, &request)
1203            .await?;
1204        let wrote = conn.last_request_bytes;
1205        drop(conn);
1206
1207        let error_code = response
1208            .topics
1209            .first()
1210            .and_then(|t| t.partitions.first())
1211            .map(|p| p.error_code)
1212            .context("OffsetCommit response named no partition")?;
1213
1214        if error_code != 0 {
1215            bail!(
1216                "commit of {topic}/{partition} offset {offset} was refused: {error_code} ({})",
1217                kafka_error_name(error_code)
1218            );
1219        }
1220
1221        info!(
1222            "Kafka client {} committed {}/{} offset {} for group {}",
1223            self.client_id, topic, partition, offset, group_id
1224        );
1225        let _ = self.status_tx.send(format!(
1226            "[CLIENT] Kafka committed {topic}/{partition} offset {offset}"
1227        ));
1228        Ok(wrote)
1229    }
1230
1231    /// Poll every assigned partition until the client goes away or the connection breaks.
1232    ///
1233    /// The first round runs immediately. Sleeping first would make a consumer sit idle for a
1234    /// whole interval before its first read, and would hide records that were already there
1235    /// when it connected.
1236    async fn poll_loop(&self, interval: Duration) {
1237        let mut first = true;
1238        loop {
1239            if !first {
1240                tokio::time::sleep(interval).await;
1241            }
1242            first = false;
1243
1244            if self.app_state.get_client(self.client_id).await.is_none() {
1245                info!("Kafka client {} stopped", self.client_id);
1246                return;
1247            }
1248
1249            let assignments = self.assignments.lock().await.clone();
1250            for assignment in assignments {
1251                match self
1252                    .fetch(
1253                        &assignment.topic,
1254                        assignment.partition,
1255                        assignment.next_offset,
1256                        DEFAULT_PARTITION_MAX_BYTES,
1257                    )
1258                    .await
1259                {
1260                    Ok((Some(event), _)) => {
1261                        if self.drive(event).await {
1262                            self.close("handler asked to disconnect").await;
1263                            return;
1264                        }
1265                    }
1266                    Ok((None, _)) => {}
1267                    Err(e) => {
1268                        error!(
1269                            "Kafka client {} polling {}/{} failed: {}",
1270                            self.client_id, assignment.topic, assignment.partition, e
1271                        );
1272                        let _ = self
1273                            .status_tx
1274                            .send(format!("[CLIENT] Kafka poll failed: {e}"));
1275                        if is_fatal(&e) {
1276                            self.mark_error(&e.to_string()).await;
1277                            return;
1278                        }
1279                    }
1280                }
1281            }
1282        }
1283    }
1284
1285    async fn tracked_offset(&self, topic: &str, partition: i32) -> Option<i64> {
1286        self.assignments
1287            .lock()
1288            .await
1289            .iter()
1290            .find(|a| a.topic == topic && a.partition == partition)
1291            .map(|a| a.next_offset)
1292    }
1293
1294    /// Record where to read next, adding the partition to the poll set if it is new. A
1295    /// handler that fetches a topic it was not configured with then keeps receiving it.
1296    async fn set_tracked_offset(&self, topic: &str, partition: i32, next_offset: i64) {
1297        let mut assignments = self.assignments.lock().await;
1298        match assignments
1299            .iter_mut()
1300            .find(|a| a.topic == topic && a.partition == partition)
1301        {
1302            Some(existing) => existing.next_offset = next_offset,
1303            None => assignments.push(Assignment {
1304                topic: topic.to_string(),
1305                partition,
1306                next_offset,
1307            }),
1308        }
1309    }
1310
1311    async fn mark_error(&self, message: &str) {
1312        // The connection is unusable, so stop offering `[ send ]` on it.
1313        self.app_state.remove_client_handle(self.client_id).await;
1314        self.app_state
1315            .update_client_status(self.client_id, ClientStatus::Error(message.to_string()))
1316            .await;
1317        let _ = self.status_tx.send("__UPDATE_UI__".to_string());
1318    }
1319
1320    async fn close(&self, why: &str) {
1321        info!(
1322            "Kafka client {} disconnecting from {}: {}",
1323            self.client_id, self.remote_addr, why
1324        );
1325        // Shutting the socket down explicitly makes the broker see a clean close rather than
1326        // waiting for its own idle timeout.
1327        let _ = self.conn.lock().await.stream.shutdown().await;
1328        // Dropping the handle also closes the command channel, which ends `command_loop`.
1329        self.app_state.remove_client_handle(self.client_id).await;
1330        self.app_state
1331            .update_client_status(self.client_id, ClientStatus::Disconnected)
1332            .await;
1333        let _ = self.status_tx.send(format!(
1334            "[CLIENT] Kafka client {} disconnected",
1335            self.client_id
1336        ));
1337        let _ = self.status_tx.send("__UPDATE_UI__".to_string());
1338    }
1339}
1340
1341/// Whether an error means the connection is unusable, as opposed to one request being refused.
1342///
1343/// A refused Produce leaves the session healthy; a desynchronised correlation id or a broken
1344/// socket does not, and continuing to poll on it would spin.
1345fn is_fatal(e: &anyhow::Error) -> bool {
1346    let text = e.to_string();
1347    e.downcast_ref::<std::io::Error>().is_some()
1348        || e.chain()
1349            .any(|c| c.downcast_ref::<std::io::Error>().is_some())
1350        || text.contains("desynchronised")
1351        || text.contains("timed out")
1352        || text.contains("implausible response size")
1353}
1354
1355/// Encode one record as a v2 record batch, the format every broker since 0.11 stores.
1356fn encode_single_record(key: Option<Vec<u8>>, value: Option<Vec<u8>>) -> Result<Bytes> {
1357    let timestamp = crate::utils::clock::SystemTime::now()
1358        .duration_since(crate::utils::clock::UNIX_EPOCH)
1359        .map(|d| d.as_millis() as i64)
1360        .unwrap_or(-1);
1361
1362    let records = vec![Record {
1363        transactional: false,
1364        control: false,
1365        partition_leader_epoch: 0,
1366        producer_id: -1,
1367        producer_epoch: -1,
1368        timestamp_type: TimestampType::Creation,
1369        offset: 0,
1370        sequence: 0,
1371        timestamp,
1372        key: key.map(Bytes::from),
1373        value: value.map(Bytes::from),
1374        headers: Default::default(),
1375    }];
1376
1377    let mut buf = Vec::new();
1378    RecordBatchEncoder::encode_with_custom_compression::<_, _, Compressor>(
1379        &mut buf,
1380        &records,
1381        &RecordEncodeOptions {
1382            version: 2,
1383            compression: Compression::None,
1384        },
1385        None,
1386    )
1387    .context("encoding the record batch")?;
1388    Ok(Bytes::from(buf))
1389}
1390
1391/// Render one fetched record for a handler.
1392///
1393/// `encode_field` is the broker's own helper: printable bytes become text tagged `utf8`,
1394/// anything else becomes hex tagged `hex`. Never base64, which models cannot read.
1395fn record_to_json(record: &Record) -> serde_json::Value {
1396    let (key, key_encoding) = match record.key.as_ref() {
1397        Some(k) => {
1398            let (v, e) = encode_field(k.as_ref(), MAX_EVENT_VALUE_BYTES);
1399            (v, Some(e))
1400        }
1401        None => (serde_json::Value::Null, None),
1402    };
1403    let (value, value_encoding) = match record.value.as_ref() {
1404        Some(v) => {
1405            let (val, e) = encode_field(v.as_ref(), MAX_EVENT_VALUE_BYTES);
1406            (val, Some(e))
1407        }
1408        None => (serde_json::Value::Null, None),
1409    };
1410
1411    serde_json::json!({
1412        "offset": record.offset,
1413        "timestamp": record.timestamp,
1414        "key": key,
1415        "key_encoding": key_encoding,
1416        "value": value,
1417        "value_encoding": value_encoding,
1418    })
1419}
1420
1421/// Turn a Metadata response into the structured shape handlers see.
1422fn metadata_to_json(response: &MetadataResponse) -> serde_json::Value {
1423    let brokers: Vec<serde_json::Value> = response
1424        .brokers
1425        .iter()
1426        .map(|b| {
1427            serde_json::json!({
1428                "node_id": b.node_id.0,
1429                "host": b.host.to_string(),
1430                "port": b.port,
1431            })
1432        })
1433        .collect();
1434
1435    let topics: Vec<serde_json::Value> = response
1436        .topics
1437        .iter()
1438        .map(|t| {
1439            serde_json::json!({
1440                "name": t.name.as_ref().map(|n| n.to_string()),
1441                "error_code": t.error_code,
1442                "error_name": kafka_error_name(t.error_code),
1443                "partitions": t.partitions.iter().map(|p| serde_json::json!({
1444                    "partition": p.partition_index,
1445                    "leader": p.leader_id.0,
1446                    "error_code": p.error_code,
1447                    "replicas": p.replica_nodes.iter().map(|r| r.0).collect::<Vec<_>>(),
1448                })).collect::<Vec<_>>(),
1449            })
1450        })
1451        .collect();
1452
1453    serde_json::json!({
1454        "cluster_id": response.cluster_id.as_ref().map(|c| c.to_string()),
1455        "controller_id": response.controller_id.0,
1456        "brokers": brokers,
1457        "topics": topics,
1458    })
1459}
1460
1461fn str_field(data: &serde_json::Value, key: &str) -> Result<String> {
1462    data.get(key)
1463        .and_then(|v| v.as_str())
1464        .map(|s| s.to_string())
1465        .with_context(|| format!("action data is missing '{key}'"))
1466}
1467
1468fn i64_field(data: &serde_json::Value, key: &str) -> Result<i64> {
1469    data.get(key)
1470        .and_then(|v| v.as_i64())
1471        .with_context(|| format!("action data is missing '{key}'"))
1472}