Skip to main content

netget/client/ftp/
mod.rs

1//! FTP client implementation
2//!
3//! LLM-controlled FTP client that connects to remote FTP servers
4//! and allows the LLM to send commands and interpret responses.
5
6pub mod actions;
7
8pub use actions::FtpClientProtocol;
9
10use anyhow::{Context, Result};
11use std::net::SocketAddr;
12use std::sync::Arc;
13use tokio::io::{AsyncWriteExt, BufReader};
14use tokio::net::TcpStream;
15use tokio::sync::{mpsc, Mutex};
16use tracing::{debug, error, info, trace};
17
18use crate::client::ftp::actions::{FTP_CLIENT_CONNECTED_EVENT, FTP_CLIENT_RESPONSE_EVENT};
19use crate::client::llm_budget::call_llm_for_client;
20use crate::llm::ollama_client::OllamaClient;
21use crate::llm::ClientLlmResult;
22use crate::protocol::Event;
23use crate::state::app_state::AppState;
24use crate::state::{ClientId, ClientStatus};
25
26/// Connection state for LLM processing
27#[derive(Debug, Clone, PartialEq)]
28enum ConnectionState {
29    Idle,
30    Processing,
31    Accumulating,
32}
33
34/// Per-client data for LLM handling
35struct ClientData {
36    state: ConnectionState,
37    queued_responses: Vec<String>,
38    memory: String,
39}
40
41/// FTP client that connects to remote FTP servers
42pub struct FtpClient;
43
44impl FtpClient {
45    /// Connect to an FTP server with integrated LLM actions
46    pub async fn connect_with_llm_actions(
47        remote_addr: String,
48        llm_client: OllamaClient,
49        app_state: Arc<AppState>,
50        status_tx: mpsc::UnboundedSender<String>,
51        client_id: ClientId,
52    ) -> Result<SocketAddr> {
53        // Resolve and connect
54        let stream = TcpStream::connect(&remote_addr).await.context(format!(
55            "Failed to connect to FTP server at {}",
56            remote_addr
57        ))?;
58
59        let local_addr = stream.local_addr()?;
60        let remote_sock_addr = stream.peer_addr()?;
61
62        info!(
63            "FTP client {} connected to {} (local: {})",
64            client_id, remote_sock_addr, local_addr
65        );
66
67        // Update client state
68        app_state
69            .update_client_status(client_id, ClientStatus::Connected)
70            .await;
71        let _ = status_tx.send(format!("[CLIENT] FTP client {} connected", client_id));
72        let _ = status_tx.send("__UPDATE_UI__".to_string());
73
74        // Split stream
75        let (read_half, write_half) = tokio::io::split(stream);
76        let write_half_arc = Arc::new(Mutex::new(write_half));
77        let write_half_for_connected = write_half_arc.clone();
78
79        // Initialize client data
80        let client_data = Arc::new(Mutex::new(ClientData {
81            state: ConnectionState::Idle,
82            queued_responses: Vec::new(),
83            memory: String::new(),
84        }));
85
86        // Command channel for the dashboard's [ send_ftp_command ] / [ disconnect ] rows.
87        // Registered BEFORE the ftp_connected call: a manual handler can park that call for
88        // minutes, and an injected send must work for the whole park. The read loop below
89        // uses `read_line`, which is not cancellation-safe, so the channel is drained by its
90        // own task sharing the write half rather than by a `select!` arm.
91        let mut command_rx =
92            crate::client::command_support::register_command_channel(&app_state, client_id).await;
93        let cmd_state = app_state.clone();
94        let cmd_tx = status_tx.clone();
95        let cmd_write = write_half_arc.clone();
96        let cmd_task = tokio::spawn(async move {
97            while let Some(cmd) = command_rx.recv().await {
98                let disconnect = crate::client::command_support::handle_stream_client_command(
99                    &crate::client::ftp::actions::FtpClientProtocol,
100                    &cmd_write,
101                    cmd,
102                    client_id,
103                    &cmd_state,
104                    &cmd_tx,
105                )
106                .await;
107                if disconnect {
108                    // Half-close: the server reads EOF and closes; the read loop then sees
109                    // 0 and runs its normal disconnect path.
110                    let _ = cmd_write.lock().await.shutdown().await;
111                    break;
112                }
113            }
114        });
115        app_state.register_client_task(client_id, cmd_task).await;
116
117        // Call LLM with ftp_connected event
118        if let Some(instruction) = app_state.get_instruction_for_client(client_id).await {
119            let event = Event::new(
120                &FTP_CLIENT_CONNECTED_EVENT,
121                serde_json::json!({
122                    "remote_addr": remote_sock_addr.to_string(),
123                }),
124            );
125
126            match call_llm_for_client(
127                &llm_client,
128                &app_state,
129                client_id.to_string(),
130                &instruction,
131                &client_data.lock().await.memory,
132                Some(&event),
133                &crate::client::ftp::actions::FtpClientProtocol,
134                &status_tx,
135            )
136            .await
137            {
138                Ok(result) => {
139                    // Update memory if provided
140                    if let Some(new_memory) = result.memory_updates {
141                        client_data.lock().await.memory = new_memory;
142                    }
143
144                    // Execute actions from LLM response
145                    use crate::llm::actions::client_trait::Client;
146                    let protocol = crate::client::ftp::actions::FtpClientProtocol::new();
147                    for action in result.actions {
148                        match protocol.execute_action(action) {
149                            Ok(
150                                crate::llm::actions::client_trait::ClientActionResult::SendData(
151                                    bytes,
152                                ),
153                            ) => {
154                                let mut write_guard = write_half_for_connected.lock().await;
155                                if let Err(e) = write_guard.write_all(&bytes).await {
156                                    error!("Failed to send FTP command after connect: {}", e);
157                                } else if let Err(e) = write_guard.flush().await {
158                                    error!("Failed to flush after connect: {}", e);
159                                } else {
160                                    let cmd = String::from_utf8_lossy(&bytes);
161                                    info!("FTP client sent: {}", cmd.trim());
162                                }
163                            }
164                            Ok(
165                                crate::llm::actions::client_trait::ClientActionResult::Disconnect,
166                            ) => {
167                                info!("LLM requested disconnect after connect");
168                                app_state.remove_client_handle(client_id).await;
169                                return Ok(local_addr);
170                            }
171                            Ok(
172                                crate::llm::actions::client_trait::ClientActionResult::WaitForMore,
173                            ) => {
174                                // Just wait for data
175                            }
176                            Ok(_) => {
177                                // Other action results
178                            }
179                            Err(e) => {
180                                error!("Failed to execute FTP action after connect: {}", e);
181                            }
182                        }
183                    }
184                }
185                Err(e) => {
186                    error!("LLM error on ftp_connected event: {}", e);
187                }
188            }
189        }
190
191        // Spawn read loop
192        // Registered with AppState so stop_client can abort this task —
193        // dropping a JoinHandle only detaches it in Tokio.
194        let task_registrar = app_state.clone();
195        let task_handle = tokio::spawn(async move {
196            info!("FTP client {} read loop started", client_id);
197            let mut reader = BufReader::new(read_half);
198            let mut line = String::new();
199
200            loop {
201                line.clear();
202                match crate::client::response_reader::read_response_line(&mut reader, &mut line)
203                    .await
204                {
205                    Ok(0) => {
206                        info!("FTP client {} disconnected", client_id);
207                        app_state
208                            .update_client_status(client_id, ClientStatus::Disconnected)
209                            .await;
210                        let _ = status_tx
211                            .send(format!("[CLIENT] FTP client {} disconnected", client_id));
212                        let _ = status_tx.send("__UPDATE_UI__".to_string());
213                        break;
214                    }
215                    Ok(_) => {
216                        let response = line.trim().to_string();
217                        debug!("FTP client {} received: {}", client_id, response);
218                        trace!("FTP client {} received response: {}", client_id, response);
219
220                        // Handle response with LLM
221                        let mut client_data_lock = client_data.lock().await;
222
223                        match client_data_lock.state {
224                            ConnectionState::Idle => {
225                                // Process immediately
226                                client_data_lock.state = ConnectionState::Processing;
227                                drop(client_data_lock);
228
229                                // Call LLM
230                                if let Some(instruction) =
231                                    app_state.get_instruction_for_client(client_id).await
232                                {
233                                    let protocol = Arc::new(
234                                        crate::client::ftp::actions::FtpClientProtocol::new(),
235                                    );
236
237                                    // Parse response code if present
238                                    let response_code = response
239                                        .split_whitespace()
240                                        .next()
241                                        .and_then(|s| s.trim_end_matches('-').parse::<u16>().ok());
242
243                                    let event = Event::new(
244                                        &FTP_CLIENT_RESPONSE_EVENT,
245                                        serde_json::json!({
246                                            "response": response,
247                                            "response_code": response_code,
248                                        }),
249                                    );
250
251                                    match call_llm_for_client(
252                                        &llm_client,
253                                        &app_state,
254                                        client_id.to_string(),
255                                        &instruction,
256                                        &client_data.lock().await.memory,
257                                        Some(&event),
258                                        protocol.as_ref(),
259                                        &status_tx,
260                                    )
261                                    .await
262                                    {
263                                        Ok(ClientLlmResult {
264                                            actions,
265                                            memory_updates,
266                                        }) => {
267                                            // Update memory
268                                            if let Some(mem) = memory_updates {
269                                                client_data.lock().await.memory = mem;
270                                            }
271
272                                            // Execute actions
273                                            for action in actions {
274                                                use crate::llm::actions::client_trait::Client;
275                                                match protocol.as_ref().execute_action(action) {
276                                                    Ok(crate::llm::actions::client_trait::ClientActionResult::SendData(bytes)) => {
277                                                        let mut write_guard = write_half_arc.lock().await;
278                                                        if write_guard.write_all(&bytes).await.is_ok() {
279                                                            if write_guard.flush().await.is_ok() {
280                                                                let cmd = String::from_utf8_lossy(&bytes);
281                                                                trace!("FTP client {} sent: {}", client_id, cmd.trim());
282                                                            }
283                                                        }
284                                                    }
285                                                    Ok(crate::llm::actions::client_trait::ClientActionResult::Disconnect) => {
286                                                        info!("FTP client {} disconnecting", client_id);
287                                                        break;
288                                                    }
289                                                    _ => {}
290                                                }
291                                            }
292                                        }
293                                        Err(e) => {
294                                            error!("LLM error for FTP client {}: {}", client_id, e);
295                                        }
296                                    }
297                                }
298
299                                // Process queued responses if any
300                                let mut client_data_lock = client_data.lock().await;
301                                if !client_data_lock.queued_responses.is_empty() {
302                                    client_data_lock.queued_responses.clear();
303                                }
304                                client_data_lock.state = ConnectionState::Idle;
305                            }
306                            ConnectionState::Processing => {
307                                // Queue response
308                                client_data_lock.queued_responses.push(response);
309                                client_data_lock.state = ConnectionState::Accumulating;
310                            }
311                            ConnectionState::Accumulating => {
312                                // Continue queuing
313                                client_data_lock.queued_responses.push(response);
314                            }
315                        }
316                    }
317                    Err(e) => {
318                        error!("FTP client {} read error: {}", client_id, e);
319                        app_state
320                            .update_client_status(client_id, ClientStatus::Error(e.to_string()))
321                            .await;
322                        let _ = status_tx.send("__UPDATE_UI__".to_string());
323                        break;
324                    }
325                }
326            }
327            // Every exit path lands here: drop the command handle so the dashboard stops
328            // offering [ send ] on a dead connection (a late send then fails fast).
329            app_state.remove_client_handle(client_id).await;
330        });
331        task_registrar
332            .register_client_task(client_id, task_handle)
333            .await;
334
335        Ok(local_addr)
336    }
337}