netget/client/socket_file/mod.rs
1//! Socket File client implementation
2pub mod actions;
3
4pub use actions::SocketFileClientProtocol;
5
6use anyhow::{Context, Result};
7use std::net::SocketAddr;
8use std::sync::Arc;
9use tokio::io::{AsyncReadExt, AsyncWriteExt};
10use tokio::net::UnixStream;
11use tokio::sync::{mpsc, Mutex};
12use tracing::{error, info, trace};
13
14use crate::client::llm_budget::call_llm_for_client;
15use crate::client::socket_file::actions::SOCKET_FILE_CLIENT_DATA_RECEIVED_EVENT;
16use crate::llm::ollama_client::OllamaClient;
17use crate::llm::ClientLlmResult;
18use crate::protocol::Event;
19use crate::state::app_state::AppState;
20use crate::state::{ClientId, ClientStatus};
21
22/// Connection state for LLM processing
23#[derive(Debug, Clone, PartialEq)]
24enum ConnectionState {
25 Idle,
26 Processing,
27 Accumulating,
28}
29
30/// Per-client data for LLM handling
31struct ClientData {
32 state: ConnectionState,
33 queued_data: Vec<u8>,
34 memory: String,
35}
36
37/// Socket File client that connects to a Unix domain socket
38pub struct SocketFileClient;
39
40impl SocketFileClient {
41 /// Connect to a Unix domain socket with integrated LLM actions
42 pub async fn connect_with_llm_actions(
43 socket_path: String,
44 llm_client: OllamaClient,
45 app_state: Arc<AppState>,
46 status_tx: mpsc::UnboundedSender<String>,
47 client_id: ClientId,
48 ) -> Result<SocketAddr> {
49 // Connect to Unix domain socket
50 let stream = UnixStream::connect(&socket_path)
51 .await
52 .context(format!("Failed to connect to Unix socket {}", socket_path))?;
53
54 // Unix sockets don't have traditional socket addresses, so we create a dummy one
55 // The socket path is stored in the client's remote_addr field in app_state
56 let dummy_addr = SocketAddr::from(([127, 0, 0, 1], 0));
57
58 info!(
59 "Socket File client {} connected to {}",
60 client_id, socket_path
61 );
62
63 // Update client state
64 app_state
65 .update_client_status(client_id, ClientStatus::Connected)
66 .await;
67 let _ = status_tx.send(format!(
68 "[CLIENT] Socket File client {} connected to {}",
69 client_id, socket_path
70 ));
71 let _ = status_tx.send("__UPDATE_UI__".to_string());
72
73 // Split stream
74 let (mut read_half, write_half) = stream.into_split();
75 let write_half_arc = Arc::new(Mutex::new(write_half));
76
77 // Initialize client data
78 let client_data = Arc::new(Mutex::new(ClientData {
79 state: ConnectionState::Idle,
80 queued_data: Vec::new(),
81 memory: String::new(),
82 }));
83
84 // Spawn read loop
85 // Registered with AppState so stop_client can abort this task —
86 // dropping a JoinHandle only detaches it in Tokio.
87 let task_registrar = app_state.clone();
88 // Command channel: lets the dashboard inject actions into this loop
89 // via AppState::send_to_client (see client/command_support.rs).
90 let mut command_rx =
91 crate::client::command_support::register_command_channel(&app_state, client_id).await;
92
93 // Raise the connected event.
94 //
95 // SOCKET_FILE_CLIENT_CONNECTED_EVENT was declared and nothing raised it, so the
96 // model was only ever consulted once the peer had already said something. A client
97 // meant to speak first -- which this protocol's own example action
98 // (`send_socket_file_data`) implies -- could never do so.
99 //
100 // From a registered task, not inline: a dashboard-created client defaults to a
101 // `*` -> manual rule and awaiting a parked answer would block creation.
102 let conn_state = app_state.clone();
103 let conn_llm = llm_client.clone();
104 let conn_status = status_tx.clone();
105 let conn_write = write_half_arc.clone();
106 let conn_path = socket_path.clone();
107 let conn_task = tokio::spawn(async move {
108 let Some(instruction) = conn_state.get_instruction_for_client(client_id).await else {
109 return;
110 };
111 let protocol = crate::client::socket_file::actions::SocketFileClientProtocol::new();
112 let event = Event::new(
113 &crate::client::socket_file::actions::SOCKET_FILE_CLIENT_CONNECTED_EVENT,
114 serde_json::json!({ "socket_path": conn_path }),
115 );
116 match call_llm_for_client(
117 &conn_llm,
118 &conn_state,
119 client_id.to_string(),
120 &instruction,
121 "",
122 Some(&event),
123 &protocol,
124 &conn_status,
125 )
126 .await
127 {
128 Ok(result) => {
129 if let Some(mem) = result.memory_updates {
130 conn_state.set_memory_for_client(client_id, mem).await;
131 }
132 use crate::llm::actions::client_trait::{Client, ClientActionResult};
133 for action in result.actions {
134 if let Ok(ClientActionResult::SendData(bytes)) =
135 protocol.execute_action(action)
136 {
137 let mut guard = conn_write.lock().await;
138 if let Err(e) = guard.write_all(&bytes).await {
139 error!(
140 "Socket File client {} connect-time write failed: {}",
141 client_id, e
142 );
143 }
144 }
145 }
146 }
147 Err(e) => error!(
148 "Socket File client {} LLM error on connected event: {}",
149 client_id, e
150 ),
151 }
152 });
153 app_state.register_client_task(client_id, conn_task).await;
154
155 let task_handle = tokio::spawn(async move {
156 let mut buffer = vec![0u8; 8192];
157
158 loop {
159 let read_result = tokio::select! {
160 read = read_half.read(&mut buffer) => read,
161 Some(cmd) = command_rx.recv() => {
162 let disconnect = crate::client::command_support::handle_stream_client_command(
163 &crate::client::socket_file::actions::SocketFileClientProtocol,
164 &write_half_arc,
165 cmd,
166 client_id,
167 &app_state,
168 &status_tx,
169 )
170 .await;
171 if disconnect {
172 app_state
173 .update_client_status(client_id, ClientStatus::Disconnected)
174 .await;
175 let _ = status_tx.send(format!(
176 "[CLIENT] SocketFile client {} disconnected (injected action)",
177 client_id
178 ));
179 let _ = status_tx.send("__UPDATE_UI__".to_string());
180 break;
181 }
182 continue;
183 }
184 };
185 match read_result {
186 Ok(0) => {
187 info!("Socket File client {} disconnected", client_id);
188 app_state
189 .update_client_status(client_id, ClientStatus::Disconnected)
190 .await;
191 let _ = status_tx.send(format!(
192 "[CLIENT] Socket File client {} disconnected",
193 client_id
194 ));
195 let _ = status_tx.send("__UPDATE_UI__".to_string());
196 break;
197 }
198 Ok(n) => {
199 let data = buffer[..n].to_vec();
200 trace!("Socket File client {} received {} bytes", client_id, n);
201
202 // A round-trip is already in flight: queue these bytes and let the
203 // in-flight round-trip pick them up when it finishes. They used to be
204 // queued and then *cleared* without ever being shown to the model.
205 {
206 let mut guard = client_data.lock().await;
207 if guard.state != ConnectionState::Idle {
208 guard.queued_data.extend_from_slice(&data);
209 guard.state = ConnectionState::Accumulating;
210 continue;
211 }
212 guard.state = ConnectionState::Processing;
213 }
214
215 let mut pending = data;
216 let mut disconnect = false;
217
218 loop {
219 let Some(instruction) =
220 app_state.get_instruction_for_client(client_id).await
221 else {
222 break;
223 };
224
225 // Copy the memory OUT of the mutex before the call. Passing
226 // `&client_data.lock().await.memory` as an argument keeps the guard
227 // alive for the whole `match` that awaits it - a lock held across an
228 // LLM call, and a self-deadlock the moment the arm re-locks to store
229 // a memory update.
230 let memory = client_data.lock().await.memory.clone();
231
232 let protocol =
233 crate::client::socket_file::actions::SocketFileClientProtocol::new(
234 );
235 let (data_str, encoding) = if pending
236 .iter()
237 .all(|&b| b.is_ascii_graphic() || b.is_ascii_whitespace())
238 {
239 (String::from_utf8_lossy(&pending).to_string(), "utf8")
240 } else {
241 (hex::encode(&pending), "hex")
242 };
243 let event = Event::new(
244 &SOCKET_FILE_CLIENT_DATA_RECEIVED_EVENT,
245 serde_json::json!({
246 "data": data_str,
247 "encoding": encoding,
248 "data_length": pending.len(),
249 }),
250 );
251
252 let result = call_llm_for_client(
253 &llm_client,
254 &app_state,
255 client_id.to_string(),
256 &instruction,
257 &memory,
258 Some(&event),
259 &protocol,
260 &status_tx,
261 )
262 .await;
263
264 match result {
265 Ok(ClientLlmResult {
266 actions,
267 memory_updates,
268 }) => {
269 if let Some(mem) = memory_updates {
270 client_data.lock().await.memory = mem;
271 }
272
273 use crate::llm::actions::client_trait::{
274 Client, ClientActionResult,
275 };
276 for action in actions {
277 match protocol.execute_action(action) {
278 Ok(ClientActionResult::SendData(bytes)) => {
279 let mut guard = write_half_arc.lock().await;
280 match guard.write_all(&bytes).await {
281 Ok(()) => trace!(
282 "Socket File client {} sent {} bytes",
283 client_id,
284 bytes.len()
285 ),
286 Err(e) => error!(
287 "Socket File client {} write failed: {}",
288 client_id, e
289 ),
290 }
291 }
292 Ok(ClientActionResult::Disconnect) => {
293 // Actually hang up. This used to `break` the
294 // action loop only, so the socket stayed open and
295 // the model's decision did nothing.
296 disconnect = true;
297 break;
298 }
299 Ok(_) => {}
300 Err(e) => error!(
301 "Socket File client {} rejected an action: {}",
302 client_id, e
303 ),
304 }
305 }
306 }
307 Err(e) => {
308 error!("LLM error for Socket File client {}: {}", client_id, e);
309 }
310 }
311
312 if disconnect {
313 break;
314 }
315
316 // Anything that arrived during the call becomes the next payload,
317 // exactly as the socket_file *server* does it.
318 let queued = {
319 let mut guard = client_data.lock().await;
320 std::mem::take(&mut guard.queued_data)
321 };
322 if queued.is_empty() {
323 break;
324 }
325 pending = queued;
326 }
327
328 client_data.lock().await.state = ConnectionState::Idle;
329
330 if disconnect {
331 info!("Socket File client {} disconnecting", client_id);
332 let _ = write_half_arc.lock().await.shutdown().await;
333 app_state
334 .update_client_status(client_id, ClientStatus::Disconnected)
335 .await;
336 let _ = status_tx.send(format!(
337 "[CLIENT] Socket File client {} disconnected (model)",
338 client_id
339 ));
340 let _ = status_tx.send("__UPDATE_UI__".to_string());
341 break;
342 }
343 }
344 Err(e) => {
345 error!("Socket File client {} read error: {}", client_id, e);
346 app_state
347 .update_client_status(client_id, ClientStatus::Error(e.to_string()))
348 .await;
349 let _ = status_tx.send("__UPDATE_UI__".to_string());
350 break;
351 }
352 }
353 }
354
355 // The loop owns the only receiver; dropping the registered handle
356 // makes later send_to_client calls fail fast instead of timing out.
357 app_state.remove_client_handle(client_id).await;
358 });
359 task_registrar
360 .register_client_task(client_id, task_handle)
361 .await;
362
363 Ok(dummy_addr)
364 }
365}