netget/client/statsd/
mod.rs

1//! StatsD emitter. The command loop owns the only socket and also handles initial actions.
2pub mod actions;
3use crate::llm::actions::client_trait::{Client, ClientActionResult};
4use crate::protocol::{ConnectContext, Event};
5use crate::server::statsd::codec::{self, Dialect, DEFAULT_DIALECT};
6use crate::state::{client_handles::ClientSendOutcome, AccessLogOwner, ClientStatus};
7use crate::{console_error, console_info};
8pub use actions::StatsdClientProtocol;
9use anyhow::{bail, Context, Result};
10use serde_json::{json, Value};
11use tokio::net::UdpSocket;
12
13pub struct StatsdClient;
14impl StatsdClient {
15    pub async fn connect(ctx: ConnectContext) -> Result<std::net::SocketAddr> {
16        let dialect = Dialect::parse(
17            &ctx.startup_params
18                .as_ref()
19                .map(|p| p.get_optional_string("dialect"))
20                .transpose()?
21                .flatten()
22                .unwrap_or_else(|| DEFAULT_DIALECT.into()),
23        )?;
24        let peer = tokio::net::lookup_host(&ctx.remote_addr)
25            .await
26            .context("resolve StatsD collector")?
27            .next()
28            .context("collector resolved to no addresses")?;
29        let socket = UdpSocket::bind(if peer.is_ipv4() {
30            "0.0.0.0:0"
31        } else {
32            "[::]:0"
33        })
34        .await?;
35        socket.connect(peer).await?;
36        let local_addr = socket.local_addr()?;
37        let mut commands =
38            crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
39                .await;
40        ctx.state
41            .update_client_status(ctx.client_id, ClientStatus::Connected)
42            .await;
43        console_info!(
44            ctx.status_tx,
45            "StatsD emitter {} ready for {} ({})",
46            ctx.client_id,
47            peer,
48            dialect.as_str()
49        );
50        let registrar = ctx.state.clone();
51        let client_id = ctx.client_id;
52        let handle = tokio::spawn(async move {
53            let protocol = StatsdClientProtocol::new();
54            let connected = async {
55                let event = Event::new(
56                    &actions::STATSD_CONNECTED_EVENT,
57                    json!({"remote_addr":peer.to_string(),"local_addr":local_addr.to_string(),"dialect":dialect.as_str()}),
58                );
59                let instruction = ctx
60                    .state
61                    .get_instruction_for_client(client_id)
62                    .await
63                    .unwrap_or_default();
64                let configured = ctx
65                    .state
66                    .get_client_event_handler_config(client_id)
67                    .await
68                    .is_some_and(|c| c.find_handler("statsd_connected").is_some());
69                if instruction.is_empty() && !configured {
70                    return Ok(crate::llm::ClientLlmResult {
71                        actions: vec![],
72                        memory_updates: None,
73                    });
74                }
75                let memory = ctx
76                    .state
77                    .get_memory_for_client(client_id)
78                    .await
79                    .unwrap_or_default();
80                crate::client::llm_budget::call_llm_for_client(
81                    &ctx.llm_client,
82                    &ctx.state,
83                    client_id.to_string(),
84                    &instruction,
85                    &memory,
86                    Some(&event),
87                    &protocol,
88                    &ctx.status_tx,
89                )
90                .await
91            };
92            tokio::pin!(connected);
93            let mut connected_done = false;
94            'session: loop {
95                tokio::select! {
96                    result = &mut connected, if !connected_done => {
97                        connected_done = true;
98                        match result {
99                            Ok(result) => {
100                                if let Some(memory) = result.memory_updates { ctx.state.set_memory_for_client(client_id, memory).await; }
101                                for action in result.actions {
102                                    match Self::apply(&socket, &protocol, dialect, action).await {
103                                        Ok(ClientSendOutcome::Disconnected) => break 'session,
104                                        Ok(ClientSendOutcome::Rejected { error }) => { console_error!(ctx.status_tx, "StatsD initial action rejected: {}", error); }
105                                        Ok(_) => {},
106                                        Err(error) => { console_error!(ctx.status_tx, "StatsD send failed: {}", error); }
107                                    }
108                                }
109                            }
110                            Err(error) => { console_error!(ctx.status_tx, "StatsD connected handler failed: {}", error); }
111                        }
112                    }
113                    command = commands.recv() => {
114                        let Some(command) = command else { break; };
115                        let action = command.action.clone();
116                        let result = Self::apply(&socket, &protocol, dialect, action.clone()).await;
117                        let disconnect = matches!(result, Ok(ClientSendOutcome::Disconnected));
118                        let response = match &result { Ok(outcome) => serde_json::to_value(outcome).unwrap_or(Value::Null), Err(error) => json!({"error":error.to_string()}) };
119                        ctx.state.record_access_log(AccessLogOwner::Client(client_id.as_u32()), "StatsD", None, "injected_action", action, vec![response]).await;
120                        crate::client::command_support::reply(command, result);
121                        if disconnect { break; }
122                    }
123                }
124                let _ = ctx.status_tx.send("__UPDATE_UI__".into());
125            }
126            // Dropping the pending connected-handler future also cancels it when
127            // disconnect arrives while a manual/model response is outstanding.
128            ctx.state.remove_client_handle(client_id).await;
129            ctx.state
130                .update_client_status(client_id, ClientStatus::Disconnected)
131                .await;
132            let _ = ctx.status_tx.send("__UPDATE_UI__".into());
133        });
134        registrar.register_client_task(client_id, handle).await;
135        Ok(local_addr)
136    }
137    async fn apply(
138        socket: &UdpSocket,
139        protocol: &StatsdClientProtocol,
140        dialect: Dialect,
141        action: Value,
142    ) -> Result<ClientSendOutcome> {
143        let result = match protocol.execute_action(action) {
144            Ok(result) => result,
145            Err(error) => {
146                return Ok(ClientSendOutcome::Rejected {
147                    error: error.to_string(),
148                })
149            }
150        };
151        match result {
152            ClientActionResult::Custom { name, data } if name == "send_statsd_batch" => {
153                let records: Vec<codec::Record> = serde_json::from_value(data)?;
154                let bytes = match codec::encode_datagram(&records, dialect) {
155                    Ok(bytes) => bytes,
156                    Err(error) => {
157                        return Ok(ClientSendOutcome::Rejected {
158                            error: error.to_string(),
159                        })
160                    }
161                };
162                let bytes_sent = socket.send(&bytes).await?;
163                Ok(ClientSendOutcome::Sent { bytes_sent })
164            }
165            ClientActionResult::Disconnect => Ok(ClientSendOutcome::Disconnected),
166            _ => bail!("unsupported StatsD action result"),
167        }
168    }
169}
netget/client/statsd/mod.rs

netget/client/statsd/
mod.rs
