1pub mod actions;
59
60pub 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
98const FIRST_REQUEST_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
105
106const IDLE_BETWEEN_REQUESTS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600);
113
114const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
119
120const CONNECTION_CAP_REFUSAL: &[u8] = b"";
128
129const MIN_REQUEST_BYTES: i32 = 8;
131
132pub const MAX_REQUEST_BYTES: usize = 100 * 1024 * 1024;
136
137const MAX_TRACE_HEX_BYTES: usize = 4096;
139
140const MAX_PARTITIONS: i32 = 1024;
143
144const MAX_UNITS_PER_REQUEST: usize = 64;
148
149const MAX_RECORDS_PER_FETCH: usize = 1000;
151
152const MAX_FETCH_PAYLOAD_BYTES: usize = 8 * 1024 * 1024;
154
155const MAX_EVENT_RECORDS: usize = 20;
158const MAX_EVENT_VALUE_BYTES: usize = 1024;
159
160const 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;
166const ERR_REQUEST_TIMED_OUT: i16 = 7;
170const ERR_UNSUPPORTED_VERSION: i16 = 35;
171const ERR_UNKNOWN_TOPIC_ID: i16 = 100;
173
174pub const MAX_KNOWN_TOPIC_IDS: usize = 4096;
179
180pub const MAX_TOPIC_NAME_LEN: usize = 249;
183
184#[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 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 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 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
245pub 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
266pub 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
278pub 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
286pub struct KafkaServer {
292 cluster_id: String,
294 broker_id: i32,
296 advertised_host: String,
298 topic_ids: TopicIds,
301}
302
303impl KafkaServer {
304 fn report_topic(&self, name: &str) -> Option<uuid::Uuid> {
306 self.topic_ids.report(name)
307 }
308
309 fn topic_name(&self, id: &uuid::Uuid) -> Option<String> {
311 self.topic_ids.name(id)
312 }
313}
314
315enum Reply {
317 Data(Value),
319 Error(i16),
321 NoAnswer(i16),
329}
330
331type SharedWrite = Arc<Mutex<tokio::io::WriteHalf<TcpStream>>>;
338
339impl KafkaServer {
340 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 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 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 let task_owner = app_state.clone();
463 task_owner
464 .spawn_server_task(server_id, async move {
465 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 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 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 #[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 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 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 #[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 let read_timeout = if answered_one {
614 IDLE_BETWEEN_REQUESTS_TIMEOUT
615 } else {
616 FIRST_REQUEST_READ_TIMEOUT
617 };
618
619 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 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 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 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 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 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 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 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 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 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 {
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 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 #[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 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 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 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 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 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 #[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 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 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 #[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 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 #[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 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 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 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 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 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 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
1850struct 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
1863fn 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
1872pub(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
1897pub(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
1930fn 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}