1pub mod actions;
6pub mod wire;
7
8use crate::llm::action_helper::call_llm;
9use crate::llm::actions::protocol_trait::ActionResult;
10use crate::logging::emit::Log;
11use crate::protocol::{Event, SpawnContext};
12use crate::server::connection::ConnectionId;
13use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
14use anyhow::{Context, Result};
15use serde_json::{json, Value};
16use std::{collections::HashMap, net::SocketAddr, time::Duration};
17use tokio::{
18 io::{AsyncWriteExt, BufReader},
19 net::{TcpListener, TcpStream},
20};
21
22pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
24pub const MAX_IDLE_TIMEOUT_SECS: u64 = 3600;
25pub const MAX_LIST_OPERANDS: usize = 64;
27
28#[derive(Clone)]
29struct Config {
30 queues: Option<Vec<String>>,
31 max_job: u64,
32 idle: Duration,
33}
34
35fn config(ctx: &SpawnContext) -> Result<Config> {
36 let params = ctx.startup_params.as_ref();
37 let queues = match params
38 .map(|p| p.get_optional_array("queues"))
39 .transpose()?
40 .flatten()
41 {
42 None => None,
43 Some(list) => {
44 let mut queues = Vec::new();
45 for queue in list {
46 let queue = queue
47 .as_str()
48 .filter(|q| wire::valid_token(q))
49 .ok_or_else(|| {
50 anyhow::anyhow!("queues must be printable names without spaces")
51 })?;
52 queues.push(queue.to_string());
53 }
54 anyhow::ensure!(
55 !queues.is_empty() && queues.len() <= 256,
56 "queues must name 1 to 256 queues"
57 );
58 Some(queues)
59 }
60 };
61 let max_job = params
62 .map(|p| p.get_optional_u64("max_job_bytes"))
63 .transpose()?
64 .flatten()
65 .unwrap_or(wire::DEFAULT_MAX_JOB_BYTES);
66 anyhow::ensure!(
67 (1..=wire::MAX_JOB_BYTES_LIMIT).contains(&max_job),
68 "max_job_bytes must be between 1 and {}",
69 wire::MAX_JOB_BYTES_LIMIT
70 );
71 let idle = params
72 .map(|p| p.get_optional_u64("idle_timeout_secs"))
73 .transpose()?
74 .flatten()
75 .unwrap_or(IDLE_TIMEOUT.as_secs());
76 anyhow::ensure!(
77 (1..=MAX_IDLE_TIMEOUT_SECS).contains(&idle),
78 "idle_timeout_secs must be between 1 and {MAX_IDLE_TIMEOUT_SECS}"
79 );
80 Ok(Config {
81 queues,
82 max_job,
83 idle: Duration::from_secs(idle),
84 })
85}
86
87pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
88 let cfg = config(&ctx)?;
89 let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
90 let addr = listener.local_addr()?;
91 Log::new(Some(&ctx.status_tx)).info(format!("LPD listening on {addr}"));
92 let state = ctx.state.clone();
93 let server_id = ctx.server_id;
94 let accept = tokio::spawn(async move {
95 let limiter = crate::server::accept_bounded::ConnectionLimiter::new(
96 crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
97 );
98 loop {
99 let (socket, peer, permit) = match crate::server::accept_bounded::accept_bounded(
101 &listener,
102 &limiter,
103 b"",
104 "LPD",
105 Some(&ctx.status_tx),
106 )
107 .await
108 {
109 Ok(v) => v,
110 Err(_) => break,
111 };
112 let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
113 let now = crate::utils::clock::Instant::now();
114 ctx.state
115 .add_connection_to_server(
116 server_id,
117 ConnectionState {
118 id,
119 remote_addr: peer,
120 local_addr: addr,
121 bytes_sent: 0,
122 bytes_received: 0,
123 packets_sent: 0,
124 packets_received: 0,
125 last_activity: now,
126 status: ConnectionStatus::Active,
127 status_changed_at: now,
128 protocol_info: ProtocolConnectionInfo::empty(),
129 },
130 )
131 .await;
132 let child = ctx.clone();
133 let cfg = cfg.clone();
134 ctx.state
135 .spawn_server_task(server_id, async move {
136 let _permit = permit;
137 if let Err(e) = session(&child, id, socket, &cfg).await {
138 Log::new(Some(&child.status_tx))
139 .warn(format!("LPD connection {id} ended: {e}"));
140 }
141 child
142 .state
143 .update_connection_status(server_id, id, ConnectionStatus::Closed)
144 .await;
145 let _ = child.status_tx.send("__UPDATE_UI__".into());
146 })
147 .await;
148 }
149 });
150 state.register_server_task(server_id, accept).await;
151 Ok(addr)
152}
153
154async fn write<W: tokio::io::AsyncWrite + Unpin>(
155 ctx: &SpawnContext,
156 id: ConnectionId,
157 w: &mut W,
158 bytes: &[u8],
159) -> Result<()> {
160 tokio::time::timeout(wire::IO_TIMEOUT, w.write_all(bytes))
161 .await
162 .context("LPD write deadline")??;
163 ctx.state
164 .update_connection_stats(
165 ctx.server_id,
166 id,
167 None,
168 Some(bytes.len() as u64),
169 None,
170 Some(1),
171 )
172 .await;
173 Ok(())
174}
175
176fn outcome(ctx: &SpawnContext, id: ConnectionId, operation: &str, decision: &str) {
177 let summary = format!("LPD connection {id} operation={operation} decision={decision}");
178 let log = Log::new(Some(&ctx.status_tx));
179 if decision.starts_with("fail_closed_") || decision == "model_silent" {
180 log.error(summary);
181 } else {
182 log.info(summary);
183 }
184}
185
186async fn decision(
187 ctx: &SpawnContext,
188 id: ConnectionId,
189 event: Event,
190 expected: &str,
191) -> Result<Value> {
192 let result = match call_llm(
193 &ctx.llm_client,
194 &ctx.state,
195 ctx.server_id,
196 Some(id),
197 &event,
198 &actions::LpdProtocol,
199 )
200 .await
201 {
202 Ok(result) => result,
203 Err(error) => {
204 outcome(ctx, id, event.id(), "fail_closed_llm_error");
205 return Err(error);
206 }
207 };
208 if !result.failures.is_empty() {
209 outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
210 anyhow::bail!("LPD handler supplied an invalid action");
211 }
212 let mut found = None;
213 let mut pending = result.protocol_results;
214 while let Some(result) = pending.pop() {
215 match result {
216 ActionResult::Custom { name, data } if name == expected => {
217 if found.is_some() {
218 outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
219 anyhow::bail!("Multiple LPD answers to one event");
220 }
221 found = Some(data);
222 }
223 ActionResult::Multiple(items) => pending.extend(items),
224 _ => {}
225 }
226 }
227 if found.is_none() {
228 outcome(ctx, id, event.id(), "model_silent");
229 }
230 found.context("Handler did not answer the LPD event")
231}
232
233async fn session(
234 ctx: &SpawnContext,
235 id: ConnectionId,
236 socket: TcpStream,
237 cfg: &Config,
238) -> Result<()> {
239 let (read, mut out) = tokio::io::split(socket);
240 let mut reader = BufReader::new(read);
241 let Some(line) = wire::read_line(&mut reader, cfg.idle).await? else {
242 return Ok(());
243 };
244 ctx.state
245 .update_connection_stats(
246 ctx.server_id,
247 id,
248 Some(line.len() as u64 + 1),
249 None,
250 Some(1),
251 None,
252 )
253 .await;
254 let Some((&code, rest)) = line.split_first() else {
255 anyhow::bail!("Empty LPD command");
256 };
257 let operands = wire::operands(rest);
258 let Some(queue) = operands.first().filter(|q| wire::valid_token(q)).cloned() else {
259 anyhow::bail!("LPD command without a valid queue name");
260 };
261 let list: Vec<String> = operands.iter().skip(1).cloned().collect();
262 anyhow::ensure!(
263 list.len() <= MAX_LIST_OPERANDS + 1 && list.iter().all(|t| wire::valid_token(t)),
264 "LPD command operands out of bounds"
265 );
266 match code {
267 wire::CMD_PRINT_WAITING => {
268 Log::new(Some(&ctx.status_tx))
270 .info(format!("LPD print-waiting-jobs for queue {queue}"));
271 Ok(())
272 }
273 wire::CMD_RECEIVE_JOB => receive_job(ctx, id, &mut reader, &mut out, cfg, &queue).await,
274 wire::CMD_QUEUE_SHORT | wire::CMD_QUEUE_LONG => {
275 let long = code == wire::CMD_QUEUE_LONG;
276 let event = Event::new(
277 &actions::QUEUE_EVENT,
278 json!({"queue": queue, "long": long, "list": list}),
279 );
280 let text = match decision(ctx, id, event, "lpd_queue_status").await {
281 Ok(answer) => {
282 outcome(ctx, id, "lpd_queue_query", "model_answer");
283 let jobs = answer["jobs"].as_array().cloned().unwrap_or_default();
284 wire::render_queue(&queue, long, answer["status"].as_str(), &jobs)
285 }
286 Err(_) => format!("{queue}: queue status unavailable\n"),
287 };
288 write(ctx, id, &mut out, text.as_bytes()).await?;
289 let _ = out.shutdown().await;
290 Ok(())
291 }
292 wire::CMD_REMOVE_JOBS => {
293 let Some((agent, jobs)) = list.split_first() else {
294 anyhow::bail!("LPD remove command without an agent");
295 };
296 let event = Event::new(
297 &actions::REMOVE_EVENT,
298 json!({"queue": queue, "agent": agent, "list": jobs}),
299 );
300 let text = match decision(ctx, id, event, "lpd_remove_result").await {
302 Ok(answer) => {
303 let removed = answer["removed"].as_array().cloned().unwrap_or_default();
304 outcome(
305 ctx,
306 id,
307 "lpd_remove_request",
308 if removed.is_empty() {
309 "model_reject"
310 } else {
311 "model_answer"
312 },
313 );
314 removed
315 .iter()
316 .map(|job| match job {
317 Value::Number(n) => match n.as_u64() {
318 Some(job) => format!("job {job:03} dequeued\n"),
319 None => format!("job {n} dequeued\n"),
320 },
321 other => format!(
322 "job {} dequeued\n",
323 wire::clean(other.as_str().unwrap_or(""), 64)
324 ),
325 })
326 .collect::<String>()
327 }
328 Err(_) => String::new(),
329 };
330 write(ctx, id, &mut out, text.as_bytes()).await?;
331 let _ = out.shutdown().await;
332 Ok(())
333 }
334 other => anyhow::bail!("Unknown LPD command code {other:#04x}"),
335 }
336}
337
338async fn receive_job<R, W>(
339 ctx: &SpawnContext,
340 id: ConnectionId,
341 reader: &mut R,
342 out: &mut W,
343 cfg: &Config,
344 queue: &str,
345) -> Result<()>
346where
347 R: tokio::io::AsyncBufRead + Unpin,
348 W: tokio::io::AsyncWrite + Unpin,
349{
350 if cfg
351 .queues
352 .as_ref()
353 .is_some_and(|q| !q.iter().any(|name| name == queue))
354 {
355 Log::new(Some(&ctx.status_tx)).info(format!("LPD refused job for unknown queue {queue}"));
356 write(ctx, id, out, &[1]).await?;
357 return Ok(());
358 }
359 write(ctx, id, out, &[0]).await?;
360 let mut data: HashMap<String, Vec<u8>> = HashMap::new();
361 let mut control: Option<(String, wire::ControlFile)> = None;
362 let mut job_bytes: u64 = 0;
363 loop {
364 let Some(line) = wire::read_line(reader, cfg.idle).await? else {
365 if control.is_some() || !data.is_empty() {
366 Log::new(Some(&ctx.status_tx)).warn(format!(
367 "LPD connection {id} closed with an incomplete job; discarded"
368 ));
369 }
370 return Ok(());
371 };
372 let Some((&sub, rest)) = line.split_first() else {
373 anyhow::bail!("Empty LPD subcommand");
374 };
375 if sub == wire::SUB_ABORT {
376 data.clear();
377 control = None;
378 job_bytes = 0;
379 continue;
380 }
381 anyhow::ensure!(
382 sub == wire::SUB_CONTROL_FILE || sub == wire::SUB_DATA_FILE,
383 "Unknown LPD receive-job subcommand {sub:#04x}"
384 );
385 let operands = wire::operands(rest);
386 let count = operands.first().and_then(|c| c.parse::<u64>().ok());
387 let name = operands.get(1).filter(|n| wire::valid_token(n)).cloned();
388 let is_control = sub == wire::SUB_CONTROL_FILE;
389 let admissible = match (count, &name) {
390 (Some(count), Some(name)) if is_control => {
391 name.starts_with("cf") && count <= wire::MAX_CONTROL_FILE && control.is_none()
392 }
393 (Some(count), Some(name)) => {
394 name.starts_with("df")
395 && job_bytes.saturating_add(count) <= cfg.max_job
396 && data.len() < wire::MAX_FILES_PER_JOB
397 && !data.contains_key(name)
398 }
399 _ => false,
400 };
401 let (Some(count), Some(name), true) = (count, name, admissible) else {
402 write(ctx, id, out, &[1]).await?;
404 return Ok(());
405 };
406 write(ctx, id, out, &[0]).await?;
407 let body = wire::read_file(reader, count, cfg.idle).await?;
408 ctx.state
409 .update_connection_stats(ctx.server_id, id, Some(count + 1), None, Some(1), None)
410 .await;
411 if is_control {
412 match wire::ControlFile::parse(&body) {
413 Ok(parsed) => control = Some((name, parsed)),
414 Err(e) => {
415 Log::new(Some(&ctx.status_tx))
416 .warn(format!("LPD connection {id} refused a control file: {e}"));
417 write(ctx, id, out, &[1]).await?;
418 return Ok(());
419 }
420 }
421 } else {
422 job_bytes += count;
423 data.insert(name, body);
424 }
425 let complete = control
426 .as_ref()
427 .is_some_and(|(_, c)| c.data_files().iter().all(|n| data.contains_key(n)));
428 if !complete {
429 write(ctx, id, out, &[0]).await?;
430 continue;
431 }
432 let (control_name, parsed) = control.take().expect("complete implies a control file");
433 let accepted = decide_job(ctx, id, queue, &control_name, &parsed, &data).await;
434 write(ctx, id, out, &[if accepted { 0 } else { 1 }]).await?;
435 data.clear();
436 job_bytes = 0;
437 }
438}
439
440async fn decide_job(
441 ctx: &SpawnContext,
442 id: ConnectionId,
443 queue: &str,
444 control_name: &str,
445 control: &wire::ControlFile,
446 data: &HashMap<String, Vec<u8>>,
447) -> bool {
448 let (job_id, origin) = wire::job_number(control_name).unwrap_or_default();
449 let files: Vec<Value> = control
450 .prints
451 .iter()
452 .map(|(format, name, source)| {
453 let bytes = data.get(name).map(Vec::as_slice).unwrap_or_default();
454 let (text, _binary) = wire::preview(bytes);
455 json!({
456 "name": name,
457 "source_name": source,
458 "format": format.to_string(),
459 "size": bytes.len(),
460 "text": text,
461 })
462 })
463 .collect();
464 let event = Event::new(
465 &actions::JOB_EVENT,
466 json!({
467 "queue": queue,
468 "job_id": job_id,
469 "host": control.host.clone().or(Some(origin).filter(|o| !o.is_empty())),
470 "user": control.user,
471 "job_name": control.job_name,
472 "title": control.title,
473 "class": control.class,
474 "mail": control.mail,
475 "files": files,
476 }),
477 );
478 match decision(ctx, id, event, "lpd_job_reply").await {
479 Ok(answer) if answer["accept"] == true => {
480 outcome(ctx, id, "lpd_print_job", "model_answer");
481 true
482 }
483 Ok(_) => {
484 outcome(ctx, id, "lpd_print_job", "model_reject");
485 false
486 }
487 Err(_) => false,
488 }
489}