netget/client/mdns/mod.rs
1//! mDNS client implementation
2pub mod actions;
3
4pub use actions::MdnsClientProtocol;
5
6use anyhow::{Context, Result};
7use mdns_sd::{ServiceDaemon, ServiceEvent};
8use std::net::SocketAddr;
9use std::sync::Arc;
10use std::time::Duration;
11use tokio::sync::mpsc;
12use tracing::{error, info, trace, warn};
13
14use crate::client::llm_budget::call_llm_for_client;
15use crate::client::mdns::actions::{
16 MDNS_CLIENT_CONNECTED_EVENT, MDNS_CLIENT_SERVICE_FOUND_EVENT,
17 MDNS_CLIENT_SERVICE_RESOLVED_EVENT,
18};
19use crate::llm::ollama_client::OllamaClient;
20use crate::llm::ClientLlmResult;
21use crate::protocol::Event;
22use crate::state::app_state::AppState;
23use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
24use crate::state::{AccessLogOwner, ClientId, ClientStatus};
25
26/// What applying one executed action did. Shared by the LLM path and the command
27/// channel.
28///
29/// There is no `Sent(usize)` variant, and that is the honest shape for this client:
30/// `mdns-sd`'s `ServiceDaemon` owns the multicast socket and reports neither the
31/// query bytes it emits nor when it emits them, so no byte count we could produce
32/// would be a real one.
33enum Applied {
34 /// The action ran; the detail says what the daemon was asked to do.
35 Executed(String),
36 /// The session should end.
37 Disconnect,
38}
39
40/// mDNS client that performs service discovery on the local network
41pub struct MdnsClient;
42
43impl MdnsClient {
44 /// Initialize mDNS client with integrated LLM actions
45 pub async fn connect_with_llm_actions(
46 _remote_addr: String,
47 llm_client: OllamaClient,
48 app_state: Arc<AppState>,
49 status_tx: mpsc::UnboundedSender<String>,
50 client_id: ClientId,
51 ) -> Result<SocketAddr> {
52 info!("mDNS client {} initializing", client_id);
53
54 // Create mDNS service daemon
55 let mdns = ServiceDaemon::new().context("Failed to create mDNS service daemon")?;
56
57 // Store daemon handle in protocol_data
58 // Note: mdns daemon is not directly serializable, so we just mark it as initialized
59 app_state
60 .with_client_mut(client_id, |client| {
61 client.set_protocol_field("mdns_initialized".to_string(), serde_json::json!(true));
62 })
63 .await;
64
65 // Update status
66 app_state
67 .update_client_status(client_id, ClientStatus::Connected)
68 .await;
69 let _ = status_tx.send(format!("[CLIENT] mDNS client {} initialized", client_id));
70 let _ = status_tx.send("__UPDATE_UI__".to_string());
71
72 // Command channel: lets the dashboard (and any programmatic caller) inject
73 // actions into this client via AppState::send_to_client.
74 //
75 // Registered BEFORE the connected event is handled: a `manual` routing rule can
76 // park that event at the dashboard for minutes, and until registration the UI
77 // reports "no command channel" - reading as a protocol limitation when it is
78 // only a queue.
79 let command_rx =
80 crate::client::command_support::register_command_channel(&app_state, client_id).await;
81
82 // The command task gets a clone of the *same* daemon handle - `ServiceDaemon`
83 // is a cloneable handle onto one running daemon, so an injected browse goes out
84 // of the same sockets, to the same multicast group, as one the LLM asked for.
85 // It also keeps the daemon alive when the client has no instruction at all.
86 let cmd_daemon = mdns.clone();
87 let cmd_llm = llm_client.clone();
88 let cmd_state = app_state.clone();
89 let cmd_status = status_tx.clone();
90 let cmd_task = tokio::spawn(async move {
91 Self::command_loop(
92 command_rx, cmd_daemon, client_id, cmd_llm, cmd_state, cmd_status,
93 )
94 .await;
95 });
96 app_state.register_client_task(client_id, cmd_task).await;
97
98 // Call LLM with connected event to get initial instructions
99 if let Some(instruction) = app_state.get_instruction_for_client(client_id).await {
100 let protocol = Arc::new(crate::client::mdns::actions::MdnsClientProtocol::new());
101 let event = Event::new(
102 &MDNS_CLIENT_CONNECTED_EVENT,
103 serde_json::json!({
104 "status": "connected",
105 "message": "mDNS client ready for service discovery"
106 }),
107 );
108
109 let llm_client_clone = llm_client.clone();
110 let app_state_clone = app_state.clone();
111 let status_tx_clone = status_tx.clone();
112
113 // Registered with AppState so stop_client can abort this task —
114 // dropping a JoinHandle only detaches it in Tokio.
115 let task_registrar = app_state.clone();
116 let task_handle = tokio::spawn(async move {
117 match call_llm_for_client(
118 &llm_client_clone,
119 &app_state_clone,
120 client_id.to_string(),
121 &instruction,
122 "",
123 Some(&event),
124 protocol.as_ref(),
125 &status_tx_clone,
126 )
127 .await
128 {
129 Ok(ClientLlmResult {
130 actions,
131 memory_updates,
132 }) => {
133 // Update memory
134 if let Some(mem) = memory_updates {
135 app_state_clone.set_memory_for_client(client_id, mem).await;
136 }
137
138 // Execute initial actions
139 for action in actions {
140 if let Err(e) = Self::execute_mdns_action(
141 client_id,
142 action,
143 &mdns,
144 llm_client_clone.clone(),
145 app_state_clone.clone(),
146 status_tx_clone.clone(),
147 )
148 .await
149 {
150 error!("Failed to execute mDNS action: {}", e);
151 }
152 }
153 }
154 Err(e) => {
155 error!("LLM error for mDNS client {}: {}", client_id, e);
156 }
157 }
158 });
159 task_registrar
160 .register_client_task(client_id, task_handle)
161 .await;
162 }
163
164 // Spawn monitoring task to check for client disconnection
165 let app_state_monitor = app_state.clone();
166 // Registered with AppState so stop_client can abort this task —
167 // dropping a JoinHandle only detaches it in Tokio.
168 let task_registrar = app_state.clone();
169 let task_handle = tokio::spawn(async move {
170 loop {
171 tokio::time::sleep(Duration::from_secs(5)).await;
172
173 // Check if client was removed
174 if app_state_monitor.get_client(client_id).await.is_none() {
175 info!("mDNS client {} stopped", client_id);
176 break;
177 }
178 }
179 });
180 task_registrar
181 .register_client_task(client_id, task_handle)
182 .await;
183
184 // Return a dummy local address (mDNS is multicast UDP)
185 Ok("224.0.0.251:5353".parse().unwrap())
186 }
187
188 /// Execute an mDNS action (browse, resolve, etc.)
189 ///
190 /// Shared by the connected-event path and injected commands, so a `[ send ]` from
191 /// the dashboard drives exactly the same daemon calls the LLM would.
192 /// Run one mDNS verb against the daemon. Raises no event and calls no LLM.
193 ///
194 /// Follow-up actions the model returns for a service-found or service-resolved event
195 /// run through here. Routing them back through `execute_mdns_action` would recurse
196 /// directly -- those events are raised from inside that very function -- and would let
197 /// one discovered service spawn another browse loop, and that one another, without
198 /// bound.
199 ///
200 /// `browse_service` is therefore drained for a bounded window here rather than given
201 /// its own long-lived task: a follow-up browse is a look, not a subscription.
202 async fn run_mdns_op(
203 client_id: ClientId,
204 mdns: &ServiceDaemon,
205 name: &str,
206 data: &serde_json::Value,
207 ) -> Result<String> {
208 match name {
209 "resolve_hostname" => {
210 let hostname = data["hostname"]
211 .as_str()
212 .ok_or_else(|| anyhow::anyhow!("Missing hostname"))?;
213 let events = mdns.resolve_hostname(hostname, Some(5000))?;
214 let mut found: Vec<String> = Vec::new();
215 while let Ok(event) = events.recv_async().await {
216 match event {
217 mdns_sd::HostnameResolutionEvent::AddressesFound(_, addrs) => {
218 found.extend(addrs.iter().map(|a| a.to_string()));
219 break;
220 }
221 mdns_sd::HostnameResolutionEvent::SearchTimeout(_)
222 | mdns_sd::HostnameResolutionEvent::SearchStopped(_) => break,
223 _ => continue,
224 }
225 }
226 found.sort();
227 Ok(format!("resolve_hostname '{hostname}' -> {found:?}"))
228 }
229 "browse_service" => {
230 let service_type = data["service_type"]
231 .as_str()
232 .ok_or_else(|| anyhow::anyhow!("Missing service_type"))?;
233 let rx = mdns.browse(service_type)?;
234 let mut seen: Vec<String> = Vec::new();
235 let deadline = tokio::time::sleep(std::time::Duration::from_secs(2));
236 tokio::pin!(deadline);
237 loop {
238 tokio::select! {
239 _ = &mut deadline => break,
240 ev = rx.recv_async() => match ev {
241 Ok(mdns_sd::ServiceEvent::ServiceFound(_, full)) => seen.push(full),
242 Ok(_) => continue,
243 Err(_) => break,
244 },
245 }
246 }
247 let _ = mdns.stop_browse(service_type);
248 seen.sort();
249 seen.dedup();
250 info!("mDNS client {} follow-up browse saw {:?}", client_id, seen);
251 Ok(format!(
252 "browse_service '{service_type}' -> {} found",
253 seen.len()
254 ))
255 }
256 other => Ok(format!("'{other}' has no non-notifying path; skipped")),
257 }
258 }
259
260 /// Execute the actions the model returned for a browse event.
261 ///
262 /// They were discarded at both sites (`actions: _`), so a model that saw a service
263 /// appear and wanted to resolve its hostname was ignored -- discovery reported
264 /// nothing actionable back.
265 async fn run_browse_follow_ups(
266 client_id: ClientId,
267 actions: Vec<serde_json::Value>,
268 mdns: &ServiceDaemon,
269 protocol: &crate::client::mdns::actions::MdnsClientProtocol,
270 ) {
271 use crate::llm::actions::client_trait::{Client, ClientActionResult};
272 for action in actions {
273 let Ok(ClientActionResult::Custom { name, data }) =
274 protocol.execute_action(action.clone())
275 else {
276 continue;
277 };
278 match Self::run_mdns_op(client_id, mdns, &name, &data).await {
279 Ok(detail) => info!("mDNS client {} follow-up: {}", client_id, detail),
280 Err(e) => error!("mDNS client {} follow-up {} failed: {}", client_id, name, e),
281 }
282 }
283 }
284
285 async fn execute_mdns_action(
286 client_id: ClientId,
287 action: serde_json::Value,
288 mdns: &ServiceDaemon,
289 llm_client: OllamaClient,
290 app_state: Arc<AppState>,
291 status_tx: mpsc::UnboundedSender<String>,
292 ) -> Result<Applied> {
293 use crate::llm::actions::client_trait::Client;
294 let protocol = Arc::new(crate::client::mdns::actions::MdnsClientProtocol::new());
295
296 match protocol.as_ref().execute_action(action.clone()) {
297 Ok(crate::llm::actions::client_trait::ClientActionResult::Custom { name, data }) => {
298 match name.as_str() {
299 "browse_service" => {
300 let service_type = data["service_type"]
301 .as_str()
302 .ok_or_else(|| anyhow::anyhow!("Missing service_type"))?;
303
304 info!(
305 "mDNS client {} browsing for service: {}",
306 client_id, service_type
307 );
308 let _ = status_tx.send(format!(
309 "[CLIENT] Browsing for mDNS service: {}",
310 service_type
311 ));
312
313 // Start browsing
314 let receiver = mdns
315 .browse(service_type)
316 .context("Failed to browse service")?;
317
318 // Spawn task to handle browse events
319 let llm_client_browse = llm_client.clone();
320 let app_state_browse = app_state.clone();
321 let status_tx_browse = status_tx.clone();
322 let protocol_browse = protocol.clone();
323 // The browse task needs the daemon so a follow-up the model asks
324 // for can actually run; ServiceDaemon is a cheap handle clone.
325 let mdns_browse = mdns.clone();
326
327 // Registered with AppState so stop_client can abort this task —
328 // dropping a JoinHandle only detaches it in Tokio.
329 let task_registrar = app_state.clone();
330 let task_handle = tokio::spawn(async move {
331 loop {
332 // `recv_timeout` is a *synchronous* blocking call, so it
333 // runs on the blocking pool rather than a runtime worker.
334 // Called inline it parks the worker for up to 10s at a
335 // time, and on a current-thread runtime that stalls
336 // everything else the client is doing - including the
337 // command channel, which is how this was found.
338 let rx = receiver.clone();
339 let recv_result = match tokio::task::spawn_blocking(move || {
340 rx.recv_timeout(Duration::from_secs(10))
341 })
342 .await
343 {
344 Ok(result) => result,
345 Err(e) => {
346 error!("mDNS browse receive task failed: {}", e);
347 break;
348 }
349 };
350
351 match recv_result {
352 Ok(event) => {
353 match event {
354 ServiceEvent::ServiceFound(service_type, fullname) => {
355 trace!(
356 "mDNS service found: {} ({})",
357 fullname,
358 service_type
359 );
360
361 // Call LLM with service found event
362 if let Some(instruction) = app_state_browse
363 .get_instruction_for_client(client_id)
364 .await
365 {
366 let llm_event = Event::new(
367 &MDNS_CLIENT_SERVICE_FOUND_EVENT,
368 serde_json::json!({
369 "service_type": service_type,
370 "fullname": fullname,
371 }),
372 );
373
374 let memory = app_state_browse
375 .get_memory_for_client(client_id)
376 .await
377 .unwrap_or_default();
378
379 match call_llm_for_client(
380 &llm_client_browse,
381 &app_state_browse,
382 client_id.to_string(),
383 &instruction,
384 &memory,
385 Some(&llm_event),
386 protocol_browse.as_ref(),
387 &status_tx_browse,
388 )
389 .await
390 {
391 Ok(ClientLlmResult {
392 actions,
393 memory_updates,
394 }) => {
395 if let Some(mem) = memory_updates {
396 app_state_browse
397 .set_memory_for_client(
398 client_id, mem,
399 )
400 .await;
401 }
402 Self::run_browse_follow_ups(
403 client_id,
404 actions,
405 &mdns_browse,
406 protocol_browse.as_ref(),
407 )
408 .await;
409 }
410 Err(e) => {
411 error!("LLM error processing service found: {}", e);
412 }
413 }
414 }
415 }
416 ServiceEvent::ServiceResolved(info) => {
417 let first_addr = info
418 .get_addresses()
419 .iter()
420 .next()
421 .map(|scoped| scoped.to_string())
422 .unwrap_or_else(|| "0.0.0.0".to_string());
423
424 info!(
425 "mDNS service resolved: {} at {}:{}",
426 info.get_fullname(),
427 first_addr,
428 info.get_port()
429 );
430
431 // Call LLM with service resolved event
432 if let Some(instruction) = app_state_browse
433 .get_instruction_for_client(client_id)
434 .await
435 {
436 let llm_event = Event::new(
437 &MDNS_CLIENT_SERVICE_RESOLVED_EVENT,
438 serde_json::json!({
439 "fullname": info.get_fullname(),
440 "hostname": info.get_hostname(),
441 "addresses": info.get_addresses().iter().map(|a| a.to_string()).collect::<Vec<_>>(),
442 "port": info.get_port(),
443 "properties": info.get_properties().iter().map(|p| format!("{}={}", p.key(), p.val_str())).collect::<Vec<_>>(),
444 }),
445 );
446
447 let memory = app_state_browse
448 .get_memory_for_client(client_id)
449 .await
450 .unwrap_or_default();
451
452 match call_llm_for_client(
453 &llm_client_browse,
454 &app_state_browse,
455 client_id.to_string(),
456 &instruction,
457 &memory,
458 Some(&llm_event),
459 protocol_browse.as_ref(),
460 &status_tx_browse,
461 )
462 .await
463 {
464 Ok(ClientLlmResult {
465 actions,
466 memory_updates,
467 }) => {
468 if let Some(mem) = memory_updates {
469 app_state_browse
470 .set_memory_for_client(
471 client_id, mem,
472 )
473 .await;
474 }
475 Self::run_browse_follow_ups(
476 client_id,
477 actions,
478 &mdns_browse,
479 protocol_browse.as_ref(),
480 )
481 .await;
482 }
483 Err(e) => {
484 error!("LLM error processing service resolved: {}", e);
485 }
486 }
487 }
488 }
489 ServiceEvent::ServiceRemoved(
490 service_type,
491 fullname,
492 ) => {
493 info!(
494 "mDNS service removed: {} ({})",
495 fullname, service_type
496 );
497 }
498 ServiceEvent::SearchStarted(service_type) => {
499 trace!("mDNS search started for: {}", service_type);
500 }
501 ServiceEvent::SearchStopped(service_type) => {
502 trace!("mDNS search stopped for: {}", service_type);
503 break;
504 }
505 _ => {
506 // Handle any other event types
507 trace!("mDNS unhandled event type");
508 }
509 }
510 }
511 Err(e) => {
512 // Check if it's a timeout or disconnection
513 if format!("{:?}", e).contains("Timeout") {
514 // Timeout is expected, check if client still exists
515 if app_state_browse
516 .get_client(client_id)
517 .await
518 .is_none()
519 {
520 info!("mDNS browse task stopping (client removed)");
521 break;
522 }
523 } else {
524 info!("mDNS browse channel error: {}", e);
525 break;
526 }
527 }
528 }
529 }
530 });
531 task_registrar
532 .register_client_task(client_id, task_handle)
533 .await;
534
535 return Ok(Applied::Executed(format!(
536 "browse_service '{service_type}' started (the mdns-sd daemon \
537 owns the multicast socket, so no byte count is observable)"
538 )));
539 }
540 "resolve_hostname" => {
541 let hostname = data["hostname"]
542 .as_str()
543 .ok_or_else(|| anyhow::anyhow!("Missing hostname"))?;
544
545 info!("mDNS client {} resolving hostname: {}", client_id, hostname);
546
547 // Use mdns to resolve hostname (timeout in milliseconds)
548 match mdns.resolve_hostname(hostname, Some(5000)) {
549 Ok(events) => {
550 // `resolve_hostname` returns a Receiver of resolution
551 // EVENTS, not addresses. This used to bind it as `addrs`
552 // and log `addrs.len()` -- the channel's queue depth,
553 // which is 0 -- so the client reported "Resolved X to 0
554 // addresses" and never surfaced a single address, however
555 // well the resolution went. Drain it instead; it is a
556 // flume receiver, so `recv_async` keeps this off the
557 // runtime's worker threads.
558 let mut found: Vec<String> = Vec::new();
559 while let Ok(event) = events.recv_async().await {
560 match event {
561 mdns_sd::HostnameResolutionEvent::AddressesFound(
562 _,
563 addrs,
564 ) => {
565 found.extend(addrs.iter().map(|a| a.to_string()));
566 break;
567 }
568 mdns_sd::HostnameResolutionEvent::SearchTimeout(_)
569 | mdns_sd::HostnameResolutionEvent::SearchStopped(_) => {
570 break
571 }
572 _ => continue,
573 }
574 }
575 found.sort();
576 info!(
577 "Resolved {} to {} address(es): {:?}",
578 hostname,
579 found.len(),
580 found
581 );
582 let _ = status_tx
583 .send(format!("[CLIENT] Resolved {}: {:?}", hostname, found));
584 return Ok(Applied::Executed(if found.is_empty() {
585 format!(
586 "resolve_hostname '{hostname}': no addresses \
587 (search timed out after 5s)"
588 )
589 } else {
590 format!(
591 "resolve_hostname '{hostname}': {} -> {}",
592 found.len(),
593 found.join(", ")
594 )
595 }));
596 }
597 Err(e) => {
598 warn!("Failed to resolve {}: {}", hostname, e);
599 let _ = status_tx.send(format!(
600 "[CLIENT] Failed to resolve {}: {}",
601 hostname, e
602 ));
603 return Err(anyhow::anyhow!(
604 "resolve_hostname '{hostname}' failed: {e}"
605 ));
606 }
607 }
608 }
609 _ => {
610 warn!("Unknown mDNS action: {}", name);
611 return Ok(Applied::Executed(format!(
612 "custom result '{name}' is not an mDNS verb"
613 )));
614 }
615 }
616 }
617 Ok(crate::llm::actions::client_trait::ClientActionResult::Disconnect) => {
618 info!("mDNS client {} disconnecting", client_id);
619 app_state
620 .update_client_status(client_id, ClientStatus::Disconnected)
621 .await;
622 let _ = status_tx.send(format!("[CLIENT] mDNS client {} disconnected", client_id));
623 return Ok(Applied::Disconnect);
624 }
625 Ok(crate::llm::actions::client_trait::ClientActionResult::WaitForMore) => {
626 trace!("mDNS client {} waiting for more data", client_id);
627 return Ok(Applied::Executed("wait_for_more".to_string()));
628 }
629 Ok(other) => Ok(Applied::Executed(format!("no wire effect: {other:?}"))),
630 Err(e) => Err(e),
631 }
632 }
633
634 /// Drain injected commands until the channel closes (the client was removed) or an
635 /// injected `disconnect` ends the session.
636 ///
637 /// Bespoke rather than `command_support::handle_stream_client_command` because
638 /// every mDNS verb yields `ClientActionResult::Custom` and there is no write half
639 /// at all - the effect goes through the `ServiceDaemon` handle. The logging and
640 /// reply are byte-for-byte what the generic helper does.
641 async fn command_loop(
642 mut command_rx: mpsc::Receiver<ClientCommand>,
643 mdns: ServiceDaemon,
644 client_id: ClientId,
645 llm_client: OllamaClient,
646 app_state: Arc<AppState>,
647 status_tx: mpsc::UnboundedSender<String>,
648 ) {
649 use crate::llm::actions::client_trait::Client;
650 use crate::llm::actions::protocol_trait::Protocol;
651
652 let protocol = crate::client::mdns::actions::MdnsClientProtocol::new();
653
654 while let Some(command) = command_rx.recv().await {
655 let action = command.action.clone();
656 let outcome = Self::execute_mdns_action(
657 client_id,
658 action.clone(),
659 &mdns,
660 llm_client.clone(),
661 app_state.clone(),
662 status_tx.clone(),
663 )
664 .await
665 .map(|applied| match applied {
666 Applied::Executed(detail) => ClientSendOutcome::Executed { detail },
667 Applied::Disconnect => ClientSendOutcome::Disconnected,
668 });
669
670 // `execute_mdns_action` folds an unknown action name and a failed daemon
671 // call into the same `Err`; only the former is a rejection, so classify it
672 // here rather than reporting a bad daemon call as a bad action.
673 let outcome = match outcome {
674 Err(e) if protocol.execute_action(action.clone()).is_err() => {
675 Ok(ClientSendOutcome::Rejected {
676 error: e.to_string(),
677 })
678 }
679 other => other,
680 };
681
682 let outcome_json = match &outcome {
683 Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
684 Err(e) => serde_json::json!({"error": e.to_string()}),
685 };
686 app_state
687 .record_access_log(
688 AccessLogOwner::Client(client_id.as_u32()),
689 protocol.protocol_name(),
690 None,
691 "injected_action",
692 action,
693 vec![outcome_json],
694 )
695 .await;
696
697 let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
698 if let Err(e) = &outcome {
699 error!("mDNS client {} injected action failed: {}", client_id, e);
700 let _ = status_tx.send(format!(
701 "[WARN] Client {} injected action failed: {}",
702 client_id, e
703 ));
704 }
705 let _ = status_tx.send("__UPDATE_UI__".to_string());
706 crate::client::command_support::reply(command, outcome);
707
708 if disconnect {
709 break;
710 }
711 }
712
713 // Every exit path lands here: drop the command handle so the dashboard stops
714 // offering [ send ] on a dead client (a late send then fails fast).
715 app_state.remove_client_handle(client_id).await;
716 let _ = status_tx.send("__UPDATE_UI__".to_string());
717 }
718}