netget/client/gelf/
mod.rs

1//! GELF UDP/TCP 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::gelf::codec::{self, Compression, Transport};
6use crate::state::{client_handles::ClientSendOutcome, AccessLogOwner, ClientStatus};
7use crate::{console_error, console_info};
8pub use actions::GelfClientProtocol;
9use anyhow::{bail, Context, Result};
10use serde_json::{json, Value};
11use std::time::Duration;
12use tokio::{
13    io::{AsyncReadExt, AsyncWriteExt},
14    net::TcpStream,
15};
16const IO_TIMEOUT: Duration = Duration::from_secs(10);
17
18pub struct GelfClient;
19impl GelfClient {
20    pub async fn connect(ctx: ConnectContext) -> Result<std::net::SocketAddr> {
21        let transport = Transport::parse(
22            &ctx.startup_params
23                .as_ref()
24                .map(|p| p.get_optional_string("transport"))
25                .transpose()?
26                .flatten()
27                .unwrap_or_else(|| codec::DEFAULT_TRANSPORT.into()),
28        )?;
29        let compression = Compression::parse(
30            &ctx.startup_params
31                .as_ref()
32                .map(|p| p.get_optional_string("compression"))
33                .transpose()?
34                .flatten()
35                .unwrap_or_else(|| codec::DEFAULT_COMPRESSION.into()),
36            transport,
37        )?;
38        let chunk_size = usize::try_from(
39            ctx.startup_params
40                .as_ref()
41                .map(|p| p.get_optional_u64("chunk_size"))
42                .transpose()?
43                .flatten()
44                .unwrap_or(codec::DEFAULT_CHUNK_SIZE as u64),
45        )?;
46        anyhow::ensure!(
47            (13..=codec::MAX_DATAGRAM_BYTES).contains(&chunk_size),
48            "chunk_size must be 13..8192"
49        );
50        let (peer, local_addr, mut reader, mut writer) = match transport {
51            Transport::Tcp => {
52                let socket = tokio::time::timeout(IO_TIMEOUT, TcpStream::connect(&ctx.remote_addr))
53                    .await
54                    .context("GELF connect deadline exceeded")??;
55                let peer = socket.peer_addr()?;
56                let local = socket.local_addr()?;
57                let (reader, writer) = socket.into_split();
58                (peer, local, Some(reader), Writer::Tcp(writer))
59            }
60            Transport::Udp => {
61                let peer =
62                    tokio::time::timeout(IO_TIMEOUT, tokio::net::lookup_host(&ctx.remote_addr))
63                        .await
64                        .context("GELF resolve deadline exceeded")??
65                        .next()
66                        .context("collector resolved to no addresses")?;
67                let socket = tokio::net::UdpSocket::bind(if peer.is_ipv4() {
68                    "0.0.0.0:0"
69                } else {
70                    "[::]:0"
71                })
72                .await?;
73                socket.connect(peer).await?;
74                let local = socket.local_addr()?;
75                (peer, local, None, Writer::Udp(socket))
76            }
77        };
78        let mut commands =
79            crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
80                .await;
81        ctx.state
82            .update_client_status(ctx.client_id, ClientStatus::Connected)
83            .await;
84        console_info!(
85            ctx.status_tx,
86            "GELF emitter {} connected to {}",
87            ctx.client_id,
88            peer
89        );
90        let registrar = ctx.state.clone();
91        let client_id = ctx.client_id;
92        let handle = tokio::spawn(async move {
93            let protocol = GelfClientProtocol::new();
94            let connected = async {
95                let event = Event::new(
96                    &actions::GELF_CONNECTED_EVENT,
97                    json!({"remote_addr":peer.to_string(),"local_addr":local_addr.to_string(),"transport":transport.as_str()}),
98                );
99                let instruction = ctx
100                    .state
101                    .get_instruction_for_client(client_id)
102                    .await
103                    .unwrap_or_default();
104                let configured = ctx
105                    .state
106                    .get_client_event_handler_config(client_id)
107                    .await
108                    .is_some_and(|c| c.find_handler("gelf_connected").is_some());
109                if instruction.is_empty() && !configured {
110                    return Ok(crate::llm::ClientLlmResult {
111                        actions: vec![],
112                        memory_updates: None,
113                    });
114                }
115                let memory = ctx
116                    .state
117                    .get_memory_for_client(client_id)
118                    .await
119                    .unwrap_or_default();
120                crate::client::llm_budget::call_llm_for_client(
121                    &ctx.llm_client,
122                    &ctx.state,
123                    client_id.to_string(),
124                    &instruction,
125                    &memory,
126                    Some(&event),
127                    &protocol,
128                    &ctx.status_tx,
129                )
130                .await
131            };
132            tokio::pin!(connected);
133            let mut connected_done = false;
134            let mut unexpected = [0u8; 1];
135            'session: loop {
136                tokio::select! {
137                    result = async { match reader.as_mut() { Some(reader)=>reader.read(&mut unexpected).await, None=>std::future::pending().await } } => {
138                        match result {
139                            Ok(0) => { console_info!(ctx.status_tx, "GELF collector closed the connection"); },
140                            Ok(_) => { console_error!(ctx.status_tx, "GELF collector sent unexpected response data; closing"); },
141                            Err(error) => { console_error!(ctx.status_tx, "GELF collector read failed: {}", error); }
142                        }
143                        break;
144                    }
145                    result = &mut connected, if !connected_done => {
146                        connected_done = true;
147                        match result {
148                            Ok(result) => {
149                                if let Some(memory) = result.memory_updates { ctx.state.set_memory_for_client(client_id, memory).await; }
150                                for action in result.actions {
151                                    match Self::apply(&mut writer, &protocol, compression, chunk_size, action).await {
152                                        Ok(ClientSendOutcome::Disconnected) => break 'session,
153                                        Ok(ClientSendOutcome::Rejected { error }) => { console_error!(ctx.status_tx, "GELF initial action rejected: {}", error); }
154                                        Ok(_) => {},
155                                        Err(error) => { console_error!(ctx.status_tx, "GELF send failed: {}", error); break 'session; }
156                                    }
157                                }
158                            }
159                            Err(error) => { console_error!(ctx.status_tx, "GELF connected handler failed: {}", error); }
160                        }
161                    }
162                    command = commands.recv() => {
163                        let Some(command) = command else { break; };
164                        let action = command.action.clone();
165                        let result = Self::apply(&mut writer, &protocol, compression, chunk_size, action.clone()).await;
166                        let disconnect = result.is_err() || matches!(result, Ok(ClientSendOutcome::Disconnected));
167                        let response = match &result { Ok(outcome) => serde_json::to_value(outcome).unwrap_or(Value::Null), Err(error) => json!({"error":error.to_string()}) };
168                        ctx.state.record_access_log(AccessLogOwner::Client(client_id.as_u32()), "GELF", None, "injected_action", action, vec![response]).await;
169                        crate::client::command_support::reply(command, result);
170                        if disconnect { break; }
171                    }
172                }
173                let _ = ctx.status_tx.send("__UPDATE_UI__".into());
174            }
175            // Dropping the pending connected-handler future also cancels it when
176            // disconnect arrives while a manual/model response is outstanding.
177            ctx.state.remove_client_handle(client_id).await;
178            ctx.state
179                .update_client_status(client_id, ClientStatus::Disconnected)
180                .await;
181            let _ = ctx.status_tx.send("__UPDATE_UI__".into());
182        });
183        registrar.register_client_task(client_id, handle).await;
184        Ok(local_addr)
185    }
186    async fn apply(
187        writer: &mut Writer,
188        protocol: &GelfClientProtocol,
189        compression: Compression,
190        chunk_size: usize,
191        action: Value,
192    ) -> Result<ClientSendOutcome> {
193        let result = match protocol.execute_action(action) {
194            Ok(r) => r,
195            Err(error) => {
196                return Ok(ClientSendOutcome::Rejected {
197                    error: error.to_string(),
198                })
199            }
200        };
201        match result {
202            ClientActionResult::Custom { name, data } if name == "send_gelf_message" => {
203                let message: codec::Message = serde_json::from_value(data)?;
204                // Produce and validate every byte before the first transport write.
205                let frames = match writer {
206                    Writer::Udp(_) => {
207                        codec::encode_udp(&message, compression, chunk_size, rand::random())
208                    }
209                    Writer::Tcp(_) => codec::encode_json(&message).map(|mut b| {
210                        b.push(0);
211                        vec![b]
212                    }),
213                };
214                let frames = match frames {
215                    Ok(b) => b,
216                    Err(error) => {
217                        return Ok(ClientSendOutcome::Rejected {
218                            error: error.to_string(),
219                        })
220                    }
221                };
222                let bytes_sent = frames.iter().map(Vec::len).sum();
223                tokio::time::timeout(IO_TIMEOUT,async {
224                    for bytes in frames {
225                        match writer {
226                            Writer::Tcp(stream)=>stream.write_all(&bytes).await?,
227                            Writer::Udp(socket)=>{let n=socket.send(&bytes).await?;anyhow::ensure!(n==bytes.len(),"partial UDP write");}
228                        }
229                    }
230                    Ok::<(),anyhow::Error>(())
231                }).await.context("GELF write deadline exceeded; transport closed to avoid replay after partial send")??;
232                Ok(ClientSendOutcome::Sent { bytes_sent })
233            }
234            ClientActionResult::Disconnect => Ok(ClientSendOutcome::Disconnected),
235            _ => bail!("unsupported GELF action result"),
236        }
237    }
238}
239enum Writer {
240    Tcp(tokio::net::tcp::OwnedWriteHalf),
241    Udp(tokio::net::UdpSocket),
242}
netget/client/gelf/mod.rs

netget/client/gelf/
mod.rs
