netget-tool 0.3.4

Network protocol stubbing toolkit: ~280 protocol servers and clients whose responses come from an LLM, a script, a static rule or a person.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
//! Beanstalkd work-queue server — the model is the queue.
//!
//! The client speaks first, one command per CRLF line; `put` is followed by a body of the
//! declared length. The commands whose answer is *queue content* — put, reserve, the job
//! commands, stats, list-tubes — raise an event and the model answers with a structured action.
//! The commands whose answer is *this connection's own state* — `use`, `watch`, `ignore`,
//! `list-tube-used`, `list-tubes-watched`, `quit` — and every malformed or unknown line are
//! answered here without consulting anyone.
//!
//! Five properties worth knowing before changing the loop:
//!
//! 1. **Sizes are checked before anything is buffered.** A command line is capped at upstream's
//!    224 bytes ([`wire::MAX_LINE_BYTES`]); a `put` body at 65535 ([`wire::MAX_JOB_BYTES`]),
//!    judged from the *declared* `<bytes>` before a single body byte is read.
//! 2. **The model cannot write framing, and cannot answer the wrong command.** Its actions are
//!    rendered by [`wire`]; the loop then checks the reply fits the command
//!    ([`wire::reply_fits`]). A reply that does not fit, no reply, and a backend failure all
//!    answer `INTERNAL_ERROR` (`OUT_OF_MEMORY` — upstream's "try again later" — when the
//!    backend is at capacity), and the session continues: every one of those is a complete
//!    beanstalkd answer to one command.
//! 3. **A worker waiting in `reserve` is not idle.** When the model answers `reserve` with
//!    `wait_for_beanstalkd_job`, the server owes the worker an answer, so the idle deadline
//!    does not apply. `reserve-with-timeout` is answered `TIMED_OUT` by NetGet when its own
//!    timeout runs out; a plain `reserve` waits until a reply is written to the connection from
//!    elsewhere (the dashboard's `[ message ]`) or the worker hangs up. The connection cap is
//!    what bounds how many can wait.
//! 4. **The deadlines wrap the read, not the answer.** A command parked for a human under a
//!    `manual` rule is closed by neither `first_byte_timeout_secs` nor `idle_timeout_secs`.
//! 5. **Every close lingers** (see [`linger`]), so a refusal written just before closing reaches
//!    a peer that still has pipelined input in flight instead of being destroyed by an RST.
pub mod actions;
pub mod wire;

use crate::llm::action_helper::call_llm;
use crate::llm::actions::protocol_trait::ActionResult;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::{Event, EventType};
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use anyhow::Result;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::{mpsc, Mutex, Notify};
use wire::Command;

pub use wire::{MAX_JOB_BYTES, MAX_LINE_BYTES};

/// How long a new connection may take to send its first command.
///
/// beanstalkd is client-speaks-first and real clients send a command as soon as they connect
/// (greenstalk sends `use`/`watch` from its constructor, or nothing until the first call). 300
/// seconds is the window a `manual` rule gives a human, for a NetGet TCP client parked on its
/// operator. A listener exposed to strangers should lower it through `first_byte_timeout_secs`.
const FIRST_COMMAND_TIMEOUT: Duration = Duration::from_secs(300);

/// How long the server waits for the next command after answering one — also the bound on
/// each read of a `put` body. Upstream has no idle timeout at all; a producer that goes quiet
/// for five minutes reconnects. A worker waiting in `reserve` is exempt (see the module docs).
const IDLE_TIMEOUT: Duration = Duration::from_secs(300);

/// Concurrent connections admitted before new ones are refused — the house default. It is
/// also what bounds the number of workers that can wait in `reserve` at once, since a waiting
/// reserve has no deadline of its own.
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;

/// What a peer over [`MAX_CONNECTIONS`] reads as the answer to its first command:
/// `OUT_OF_MEMORY`, upstream's "the server cannot take this now; try again later".
const CONNECTION_CAP_REFUSAL: &[u8] = b"OUT_OF_MEMORY\r\n";

/// The peer-visible answers when the model could not or did not answer. Fixed literals, so
/// nothing derived from an error reaches the wire.
const UNAVAILABLE_REPLY: &[u8] = b"INTERNAL_ERROR\r\n";
const OVERLOADED_REPLY: &[u8] = b"OUT_OF_MEMORY\r\n";

pub struct BeanstalkdServer;

impl BeanstalkdServer {
    pub async fn spawn_with_llm_actions(
        listen_addr: SocketAddr,
        llm_client: OllamaClient,
        app_state: Arc<AppState>,
        status_tx: mpsc::UnboundedSender<String>,
        server_id: crate::state::ServerId,
        first_byte_timeout_secs: Option<u64>,
        idle_timeout_secs: Option<u64>,
    ) -> Result<SocketAddr> {
        let first_timeout = first_byte_timeout_secs
            .map(Duration::from_secs)
            .unwrap_or(FIRST_COMMAND_TIMEOUT);
        let idle_timeout = idle_timeout_secs
            .map(Duration::from_secs)
            .unwrap_or(IDLE_TIMEOUT);
        let listener = TcpListener::bind(listen_addr).await?;
        let local_addr = listener.local_addr()?;

        Log::new(Some(&status_tx)).info(format!("Beanstalkd server listening on {}", local_addr));

        let protocol = Arc::new(actions::BeanstalkdProtocol::new());
        let task_registrar = app_state.clone();
        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
        let accept_handle = tokio::spawn(async move {
            loop {
                match crate::server::accept_bounded::accept_bounded(
                    &listener,
                    &limiter,
                    CONNECTION_CAP_REFUSAL,
                    "Beanstalkd",
                    Some(&status_tx),
                )
                .await
                {
                    Ok((socket, peer_addr, permit)) => {
                        let connection_id =
                            ConnectionId::new(app_state.get_next_unified_id().await);
                        use crate::state::server::{
                            ConnectionState as ServerConnectionState, ConnectionStatus,
                            ProtocolConnectionInfo,
                        };
                        let now = crate::utils::clock::Instant::now();
                        app_state
                            .add_connection_to_server(
                                server_id,
                                ServerConnectionState {
                                    id: connection_id,
                                    remote_addr: peer_addr,
                                    local_addr,
                                    bytes_sent: 0,
                                    bytes_received: 0,
                                    packets_sent: 0,
                                    packets_received: 0,
                                    last_activity: now,
                                    status: ConnectionStatus::Active,
                                    status_changed_at: now,
                                    protocol_info: ProtocolConnectionInfo::empty(),
                                },
                            )
                            .await;
                        let _ = status_tx.send("__UPDATE_UI__".to_string());
                        Log::new(Some(&status_tx))
                            .info(format!("Beanstalkd client connected from {}", peer_addr));

                        let session = Session {
                            peer_addr,
                            llm_client: llm_client.clone(),
                            app_state: app_state.clone(),
                            status_tx: status_tx.clone(),
                            server_id,
                            protocol: protocol.clone(),
                            connection_id,
                            first_timeout,
                            idle_timeout,
                        };
                        // Tracked, not detached: stop_server must abort this task too.
                        let task_owner = app_state.clone();
                        task_owner
                            .spawn_server_task(server_id, async move {
                                // Released when the connection ends; this task is the whole
                                // connection, so MAX_CONNECTIONS caps live connections.
                                let _permit = permit;
                                session.run(socket).await
                            })
                            .await;
                    }
                    Err(e) => {
                        Log::new(Some(&status_tx)).error(format!("Beanstalkd accept error: {}", e));
                        break;
                    }
                }
            }
        });

        task_registrar
            .register_server_task(server_id, accept_handle)
            .await;
        Ok(local_addr)
    }
}

/// Counts writes to one connection and wakes whoever is waiting on them.
///
/// A waiting `reserve` must end when *anything* answers it — including a reply injected from
/// the dashboard, which is written by `peer_support`'s own task and never passes through the
/// session loop. Watching the writer is the one place both paths meet.
#[derive(Default)]
struct WriteWatch {
    writes: AtomicU64,
    notify: Notify,
}

impl WriteWatch {
    fn count(&self) -> u64 {
        self.writes.load(Ordering::SeqCst)
    }
}

/// The connection's write half, reporting every non-empty write to its [`WriteWatch`].
struct WatchedWriter<W> {
    inner: W,
    watch: Arc<WriteWatch>,
}

impl<W: AsyncWrite + Unpin> AsyncWrite for WatchedWriter<W> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        let polled = Pin::new(&mut self.inner).poll_write(cx, buf);
        if let Poll::Ready(Ok(n)) = &polled {
            if *n > 0 {
                self.watch.writes.fetch_add(1, Ordering::SeqCst);
                self.watch.notify.notify_waiters();
            }
        }
        polled
    }
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }
    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }
}

type Writer = Arc<Mutex<WatchedWriter<tokio::io::WriteHalf<tokio::net::TcpStream>>>>;

struct Session {
    peer_addr: SocketAddr,
    llm_client: OllamaClient,
    app_state: Arc<AppState>,
    status_tx: mpsc::UnboundedSender<String>,
    server_id: crate::state::ServerId,
    protocol: Arc<actions::BeanstalkdProtocol>,
    connection_id: ConnectionId,
    first_timeout: Duration,
    idle_timeout: Duration,
}

/// The connection's own tube state: the one piece of state NetGet keeps, because it is about
/// this connection rather than about the queue.
struct Tubes {
    used: String,
    watched: Vec<String>,
}

/// What handling one command decided.
enum Next {
    Continue,
    /// The model left a reserve waiting for a job.
    Wait,
    Close,
}

impl Session {
    async fn run(self, socket: tokio::net::TcpStream) {
        let (mut reader, write_half) = tokio::io::split(socket);
        let watch = Arc::new(WriteWatch::default());
        let writer: Writer = Arc::new(Mutex::new(WatchedWriter {
            inner: write_half,
            watch: watch.clone(),
        }));

        // Registered before the first command, so the operator can reach this connection
        // while a command is parked for them or a worker waits in reserve.
        let peer_rx = crate::server::peer_support::register_peer_channel(
            &self.app_state,
            self.server_id,
            self.connection_id.as_u32(),
        )
        .await;
        crate::server::peer_support::spawn_peer_command_task(
            peer_rx,
            self.protocol.clone(),
            self.app_state.clone(),
            self.server_id,
            self.connection_id.as_u32(),
            writer.clone(),
            self.status_tx.clone(),
        );

        {
            let mut framer = Framer::new(&mut reader);
            self.session(&mut framer, &writer, &watch).await;
        }

        self.app_state
            .remove_peer_handle(self.server_id, self.connection_id.as_u32())
            .await;
        let _ = writer.lock().await.shutdown().await;
        linger(&mut reader).await;
        self.app_state
            .update_connection_status(
                self.server_id,
                self.connection_id,
                crate::state::server::ConnectionStatus::Closed,
            )
            .await;
        let _ = self.status_tx.send("__UPDATE_UI__".to_string());
    }

    async fn write(&self, writer: &Writer, data: &[u8]) -> std::io::Result<()> {
        {
            let mut w = writer.lock().await;
            w.write_all(data).await?;
            w.flush().await?;
        }
        self.record_sent(data.len()).await;
        Ok(())
    }

    async fn record_sent(&self, n: usize) {
        self.app_state
            .update_connection_stats(
                self.server_id,
                self.connection_id,
                None,
                Some(n as u64),
                None,
                Some(1),
            )
            .await;
    }

    async fn record_received(&self, n: usize) {
        self.app_state
            .update_connection_stats(
                self.server_id,
                self.connection_id,
                Some(n as u64),
                None,
                Some(1),
                None,
            )
            .await;
    }

    async fn session<R: AsyncRead + Unpin>(
        &self,
        framer: &mut Framer<'_, R>,
        writer: &Writer,
        watch: &Arc<WriteWatch>,
    ) {
        let log = Log::new(Some(&self.status_tx));
        let mut tubes = Tubes {
            used: "default".to_string(),
            watched: vec!["default".to_string()],
        };
        let mut answered_one = false;

        loop {
            let timeout = if answered_one {
                self.idle_timeout
            } else {
                self.first_timeout
            };
            let (line, n) = match framer.next_line(timeout).await {
                LineRead::Line(line, n) => (line, n),
                LineRead::Eof => {
                    log.info(format!("Beanstalkd client {} disconnected", self.peer_addr));
                    return;
                }
                LineRead::TimedOut => {
                    log.info(format!(
                        "Beanstalkd client {} sent no command within {}s; closing",
                        self.peer_addr,
                        timeout.as_secs()
                    ));
                    return;
                }
                LineRead::TooLong => {
                    log.warn(format!(
                        "Beanstalkd command from {} exceeded {} bytes \
                         decision=fail_closed_line_too_long",
                        self.peer_addr, MAX_LINE_BYTES
                    ));
                    let _ = self.write(writer, b"BAD_FORMAT\r\n").await;
                    return;
                }
                LineRead::Failed(e) => {
                    log.error(format!(
                        "Beanstalkd read error from {}: {}",
                        self.peer_addr, e
                    ));
                    return;
                }
            };
            self.record_received(n).await;
            log.trace(format!(
                "Beanstalkd command from {}: {:?}",
                self.peer_addr, line
            ));
            answered_one = true;
            let command = wire::parse_command(&line);

            // Commands about this connection's own state, and every malformed line.
            let fixed: Option<Vec<u8>> = match &command {
                Command::Use(tube) => {
                    tubes.used = tube.clone();
                    Some(format!("USING {tube}\r\n").into_bytes())
                }
                Command::Watch(tube) => {
                    if !tubes.watched.contains(tube) {
                        tubes.watched.push(tube.clone());
                    }
                    Some(format!("WATCHING {}\r\n", tubes.watched.len()).into_bytes())
                }
                Command::Ignore(tube) => {
                    if tubes.watched.len() == 1 && tubes.watched[0] == *tube {
                        Some(b"NOT_IGNORED\r\n".to_vec())
                    } else {
                        tubes.watched.retain(|t| t != tube);
                        Some(format!("WATCHING {}\r\n", tubes.watched.len()).into_bytes())
                    }
                }
                Command::ListTubeUsed => Some(format!("USING {}\r\n", tubes.used).into_bytes()),
                Command::ListTubesWatched => Some(
                    wire::render_tube_list(&tubes.watched)
                        .map(String::into_bytes)
                        .unwrap_or_else(|_| UNAVAILABLE_REPLY.to_vec()),
                ),
                Command::Quit => return,
                Command::BadFormat => Some(b"BAD_FORMAT\r\n".to_vec()),
                Command::Unknown => Some(b"UNKNOWN_COMMAND\r\n".to_vec()),
                _ => None,
            };
            if let Some(reply) = fixed {
                if self.write(writer, &reply).await.is_err() {
                    return;
                }
                continue;
            }

            // put: the body, judged by its declared size before any of it is read.
            let mut body: Option<Vec<u8>> = None;
            if let Command::Put { bytes, .. } = &command {
                let declared = *bytes;
                match declared {
                    Some(n) if n <= MAX_JOB_BYTES as u64 => {
                        match framer.read_exact(n as usize + 2, self.idle_timeout).await {
                            Ok(mut data) => {
                                self.record_received(data.len()).await;
                                if !data.ends_with(b"\r\n") {
                                    log.warn(format!(
                                        "Beanstalkd put from {}: body not followed by CRLF \
                                         decision=fail_closed_expected_crlf",
                                        self.peer_addr
                                    ));
                                    if self.write(writer, b"EXPECTED_CRLF\r\n").await.is_err() {
                                        return;
                                    }
                                    continue;
                                }
                                data.truncate(n as usize);
                                body = Some(data);
                            }
                            Err(reason) => {
                                log.info(format!(
                                    "Beanstalkd put from {}: body never arrived ({reason}); \
                                     closing",
                                    self.peer_addr
                                ));
                                return;
                            }
                        }
                    }
                    _ => {
                        log.warn(format!(
                            "Beanstalkd put from {} declared {} bytes, over max-job-size {} \
                             decision=fail_closed_job_too_big",
                            self.peer_addr,
                            declared
                                .map(|n| n.to_string())
                                .unwrap_or_else(|| "more than 2^64".to_string()),
                            MAX_JOB_BYTES
                        ));
                        if self.write(writer, b"JOB_TOO_BIG\r\n").await.is_err() {
                            return;
                        }
                        // Stay in step with a client that sends the body anyway, as upstream
                        // does — unless the body is so large that skipping it is not worth it.
                        match declared.map(|n| n.saturating_add(2)) {
                            Some(skip) if skip <= wire::MAX_DISCARD_BYTES => {
                                if !framer.discard(skip, self.idle_timeout).await {
                                    return;
                                }
                                continue;
                            }
                            _ => return,
                        }
                    }
                }
            }

            match self
                .answer_with_model(&command, body.as_deref(), &tubes, writer)
                .await
            {
                Next::Continue => {}
                Next::Close => return,
                Next::Wait => {
                    if !self.wait_in_reserve(&command, framer, writer, watch).await {
                        return;
                    }
                }
            }
        }
    }

    /// Hold a worker whose reserve the model left waiting. Returns `false` when the connection
    /// should close.
    async fn wait_in_reserve<R: AsyncRead + Unpin>(
        &self,
        command: &Command,
        framer: &mut Framer<'_, R>,
        writer: &Writer,
        watch: &Arc<WriteWatch>,
    ) -> bool {
        let log = Log::new(Some(&self.status_tx));
        let deadline = match command {
            Command::ReserveWithTimeout(secs) => {
                Some(tokio::time::Instant::now() + Duration::from_secs(u64::from(*secs)))
            }
            _ => None,
        };
        let start = watch.count();
        loop {
            let notified = watch.notify.notified();
            tokio::pin!(notified);
            notified.as_mut().enable();
            if watch.count() != start {
                log.info(format!(
                    "Beanstalkd reserve from {} answered by a reply written to the connection",
                    self.peer_addr
                ));
                return true;
            }
            let timer = async {
                match deadline {
                    Some(at) => tokio::time::sleep_until(at).await,
                    None => std::future::pending::<()>().await,
                }
            };
            tokio::select! {
                read = framer.fill() => match read {
                    Ok(0) => {
                        log.info(format!(
                            "Beanstalkd client {} hung up while waiting in reserve",
                            self.peer_addr
                        ));
                        return false;
                    }
                    Ok(n) => {
                        log.warn(format!(
                            "Beanstalkd client {} sent {n} bytes while its reserve was waiting; \
                             the reserve is abandoned and the command is read",
                            self.peer_addr
                        ));
                        return true;
                    }
                    Err(e) => {
                        log.error(format!("Beanstalkd read error from {}: {}", self.peer_addr, e));
                        return false;
                    }
                },
                _ = timer => {
                    // Written only if nothing answered the reserve meanwhile — checked under
                    // the writer lock, which an injected reply also takes.
                    let wrote = {
                        let mut w = writer.lock().await;
                        if watch.count() != start {
                            false
                        } else {
                            if w.write_all(b"TIMED_OUT\r\n").await.is_err()
                                || w.flush().await.is_err()
                            {
                                return false;
                            }
                            true
                        }
                    };
                    if wrote {
                        self.record_sent(11).await;
                        log.info(format!(
                            "Beanstalkd reserve from {} decision=model_wait_timed_out",
                            self.peer_addr
                        ));
                    }
                    return true;
                }
                _ = &mut notified => continue,
            }
        }
    }

    /// Raise the command's event and write what the model decided.
    async fn answer_with_model(
        &self,
        command: &Command,
        body: Option<&[u8]>,
        tubes: &Tubes,
        writer: &Writer,
    ) -> Next {
        let log = Log::new(Some(&self.status_tx));
        let Some((event_type, data)) = event_for(command, body, tubes) else {
            return Next::Continue;
        };
        let event = Event::new(event_type, data);

        let result = match call_llm(
            &self.llm_client,
            &self.app_state,
            self.server_id,
            Some(self.connection_id),
            &event,
            self.protocol.as_ref(),
        )
        .await
        {
            Ok(result) => result,
            Err(e) => {
                let (category, reply) = match crate::utils::WireFailure::classify(&e) {
                    crate::utils::WireFailure::Overloaded => ("overloaded", OVERLOADED_REPLY),
                    crate::utils::WireFailure::Unavailable => ("unavailable", UNAVAILABLE_REPLY),
                };
                log.warn(format!(
                    "Beanstalkd {:?} from {} decision=fail_closed_llm_error category={}",
                    command, self.peer_addr, category
                ));
                log.debug(format!("Beanstalkd LLM call failed: {}", e));
                return match self.write(writer, reply).await {
                    Ok(()) => Next::Continue,
                    Err(_) => Next::Close,
                };
            }
        };
        for message in &result.messages {
            log.info(message);
        }

        let mut replies: Vec<Vec<u8>> = Vec::new();
        let mut close = false;
        let mut wait = false;
        let mut stack = result.protocol_results;
        stack.reverse();
        while let Some(item) = stack.pop() {
            match item {
                ActionResult::Output(bytes) => replies.push(bytes),
                ActionResult::CloseConnection => close = true,
                ActionResult::WaitForMore => wait = true,
                ActionResult::Multiple(items) => stack.extend(items.into_iter().rev()),
                _ => {}
            }
        }

        let is_reserve = matches!(command, Command::Reserve | Command::ReserveWithTimeout(_));
        let Some(reply) = replies.first() else {
            if wait && is_reserve {
                log.info(format!(
                    "Beanstalkd {:?} from {} decision=model_wait",
                    command, self.peer_addr
                ));
                return if close { Next::Close } else { Next::Wait };
            }
            let decision = if wait {
                "fail_closed_mismatched_reply"
            } else {
                "model_silent"
            };
            log.warn(format!(
                "Beanstalkd {:?} from {} decision={} ({} failed action(s)); answering \
                 INTERNAL_ERROR",
                command,
                self.peer_addr,
                decision,
                result.failures.len()
            ));
            return match self.write(writer, UNAVAILABLE_REPLY).await {
                Ok(()) if !close => Next::Continue,
                _ => Next::Close,
            };
        };

        if !wire::reply_fits(command, reply) {
            let (word, _) = wire::reply_head(reply);
            log.warn(format!(
                "Beanstalkd {:?} from {} decision=fail_closed_mismatched_reply reply={}; \
                 answering INTERNAL_ERROR",
                command, self.peer_addr, word
            ));
            return match self.write(writer, UNAVAILABLE_REPLY).await {
                Ok(()) if !close => Next::Continue,
                _ => Next::Close,
            };
        }
        if replies.len() > 1 {
            log.warn(format!(
                "Beanstalkd {:?}: the model produced {} replies to one command; sending the first",
                command,
                replies.len()
            ));
        }
        let (word, _) = wire::reply_head(reply);
        let decision = if wire::is_refusal(reply) {
            "model_reject"
        } else {
            "model_answer"
        };
        log.info(format!(
            "Beanstalkd {:?} from {} decision={} reply={}",
            command, self.peer_addr, decision, word
        ));
        if self.write(writer, reply).await.is_err() || close {
            return Next::Close;
        }
        Next::Continue
    }
}

/// The event a command raises and its data, for the commands the model answers.
fn event_for(
    command: &Command,
    body: Option<&[u8]>,
    tubes: &Tubes,
) -> Option<(&'static EventType, serde_json::Value)> {
    use actions::{
        BEANSTALKD_JOB_COMMAND_EVENT, BEANSTALKD_PUT_EVENT, BEANSTALKD_RESERVE_EVENT,
        BEANSTALKD_STATS_EVENT,
    };
    use serde_json::json;
    let used = &tubes.used;
    let job = |name: &str, id: Option<u64>| {
        let mut data = json!({"command": name, "tube": used});
        if let Some(id) = id {
            data["job_id"] = json!(id);
        }
        (&*BEANSTALKD_JOB_COMMAND_EVENT, data)
    };
    Some(match command {
        Command::Put {
            priority,
            delay,
            ttr,
            ..
        } => {
            let body = body.unwrap_or_default();
            (
                &*BEANSTALKD_PUT_EVENT,
                json!({
                    "tube": used,
                    "priority": priority,
                    "delay": delay,
                    "ttr": ttr,
                    "body": String::from_utf8_lossy(body),
                    "body_bytes": body.len(),
                }),
            )
        }
        Command::Reserve => (
            &*BEANSTALKD_RESERVE_EVENT,
            json!({
                "tubes": tubes.watched,
                "answer_with": actions::reserve_answer_with(&tubes.watched),
            }),
        ),
        Command::ReserveWithTimeout(secs) => (
            &*BEANSTALKD_RESERVE_EVENT,
            json!({
                "tubes": tubes.watched,
                "timeout_secs": secs,
                "answer_with": actions::reserve_answer_with(&tubes.watched),
            }),
        ),
        Command::ReserveJob(id) => job("reserve-job", Some(*id)),
        Command::Delete(id) => job("delete", Some(*id)),
        Command::Release {
            id,
            priority,
            delay,
        } => {
            let (event, mut data) = job("release", Some(*id));
            data["priority"] = json!(priority);
            data["delay"] = json!(delay);
            (event, data)
        }
        Command::Bury { id, priority } => {
            let (event, mut data) = job("bury", Some(*id));
            data["priority"] = json!(priority);
            (event, data)
        }
        Command::Touch(id) => job("touch", Some(*id)),
        Command::Peek(id) => job("peek", Some(*id)),
        Command::PeekReady => job("peek-ready", None),
        Command::PeekDelayed => job("peek-delayed", None),
        Command::PeekBuried => job("peek-buried", None),
        Command::Kick(bound) => {
            let (event, mut data) = job("kick", None);
            data["bound"] = json!(bound);
            (event, data)
        }
        Command::KickJob(id) => job("kick-job", Some(*id)),
        Command::PauseTube { tube, delay } => (
            &*BEANSTALKD_JOB_COMMAND_EVENT,
            json!({"command": "pause-tube", "tube": tube, "delay": delay}),
        ),
        Command::Stats => (
            &*BEANSTALKD_STATS_EVENT,
            json!({"scope": "server", "answer_with": actions::stats_answer_with("server")}),
        ),
        Command::StatsTube(tube) => (
            &*BEANSTALKD_STATS_EVENT,
            json!({
                "scope": "tube",
                "tube": tube,
                "answer_with": actions::stats_answer_with("tube"),
            }),
        ),
        Command::StatsJob(id) => (
            &*BEANSTALKD_STATS_EVENT,
            json!({
                "scope": "job",
                "job_id": id,
                "answer_with": actions::stats_answer_with("job"),
            }),
        ),
        Command::ListTubes => (
            &*BEANSTALKD_STATS_EVENT,
            json!({"scope": "tubes", "answer_with": actions::stats_answer_with("tubes")}),
        ),
        _ => return None,
    })
}

/// How long, and how much, the server keeps reading after it has sent its last reply and
/// half-closed. Closing a socket with unread input makes the kernel send RST instead of FIN,
/// and an RST can destroy a reply the peer has not read yet — a `BAD_FORMAT` for an over-long
/// line with pipelined commands behind it would arrive as "connection reset". Bounded both
/// ways so a peer that keeps sending cannot hold the task open.
const LINGER_TIME: Duration = Duration::from_secs(2);
const LINGER_BYTES: usize = 64 * 1024;

async fn linger<R: AsyncRead + Unpin>(reader: &mut R) {
    let deadline = tokio::time::Instant::now() + LINGER_TIME;
    let mut sink = [0u8; 4096];
    let mut drained = 0usize;
    while drained < LINGER_BYTES {
        match tokio::time::timeout_at(deadline, reader.read(&mut sink)).await {
            Ok(Ok(n)) if n > 0 => drained += n,
            _ => break,
        }
    }
}

enum LineRead {
    /// A line without its terminator, and the wire bytes it consumed.
    Line(String, usize),
    Eof,
    TimedOut,
    /// The line exceeded [`MAX_LINE_BYTES`] including its CRLF.
    TooLong,
    Failed(std::io::Error),
}

/// Accumulates reads into lines and bodies, keeping anything past what was asked for — a
/// client sends `put …\r\n<body>\r\n` in one write, and may pipeline further commands.
///
/// `pending` never holds more than one read past [`MAX_LINE_BYTES`] while a line is being
/// looked for, or one read past `MAX_JOB_BYTES + 2` while a body is.
struct Framer<'a, R> {
    reader: &'a mut R,
    pending: Vec<u8>,
    chunk: Vec<u8>,
}

impl<'a, R: AsyncRead + Unpin> Framer<'a, R> {
    fn new(reader: &'a mut R) -> Self {
        Self {
            reader,
            pending: Vec::new(),
            chunk: vec![0u8; 4096],
        }
    }

    /// One read into `pending`. `Ok(0)` is EOF. Cancel-safe: a read that has not completed
    /// has consumed nothing.
    async fn fill(&mut self) -> std::io::Result<usize> {
        let n = self.reader.read(&mut self.chunk).await?;
        self.pending.extend_from_slice(&self.chunk[..n]);
        Ok(n)
    }

    /// `timeout` bounds the wait for *more bytes*, not the whole line.
    async fn next_line(&mut self, timeout: Duration) -> LineRead {
        loop {
            if let Some(idx) = self.pending.iter().position(|b| *b == b'\n') {
                let consumed = idx + 1;
                if consumed > MAX_LINE_BYTES {
                    return LineRead::TooLong;
                }
                let line: Vec<u8> = self.pending.drain(..consumed).collect();
                let text = String::from_utf8_lossy(&line[..idx]);
                return LineRead::Line(text.trim_end_matches('\r').to_string(), consumed);
            }
            // No newline in MAX_LINE_BYTES bytes: even if the next byte is one, the line is
            // longer than upstream allows.
            if self.pending.len() >= MAX_LINE_BYTES {
                return LineRead::TooLong;
            }
            match tokio::time::timeout(timeout, self.fill()).await {
                Err(_) => return LineRead::TimedOut,
                Ok(Ok(0)) => {
                    if self.pending.is_empty() {
                        return LineRead::Eof;
                    }
                    let line = std::mem::take(&mut self.pending);
                    let consumed = line.len();
                    let text = String::from_utf8_lossy(&line);
                    return LineRead::Line(text.trim_end_matches('\r').to_string(), consumed);
                }
                Ok(Ok(_)) => {}
                Ok(Err(e)) => return LineRead::Failed(e),
            }
        }
    }

    /// Exactly `n` bytes; `timeout` bounds each wait for more.
    async fn read_exact(&mut self, n: usize, timeout: Duration) -> Result<Vec<u8>, String> {
        while self.pending.len() < n {
            match tokio::time::timeout(timeout, self.fill()).await {
                Err(_) => return Err(format!("no bytes for {}s", timeout.as_secs())),
                Ok(Ok(0)) => return Err("EOF".to_string()),
                Ok(Ok(_)) => {}
                Ok(Err(e)) => return Err(e.to_string()),
            }
        }
        Ok(self.pending.drain(..n).collect())
    }

    /// Read and drop `n` bytes without keeping them. `false` if they never came.
    async fn discard(&mut self, mut n: u64, timeout: Duration) -> bool {
        loop {
            let take = (self.pending.len() as u64).min(n) as usize;
            self.pending.drain(..take);
            n -= take as u64;
            if n == 0 {
                return true;
            }
            match tokio::time::timeout(timeout, self.fill()).await {
                Ok(Ok(k)) if k > 0 => {}
                _ => return false,
            }
        }
    }
}