1pub 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#[derive(Debug, Clone, PartialEq)]
28enum ConnectionState {
29 Idle,
30 Processing,
31 Accumulating,
32}
33
34struct ClientData {
36 state: ConnectionState,
37 queued_responses: Vec<String>,
38 memory: String,
39}
40
41pub struct FtpClient;
43
44impl FtpClient {
45 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 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 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 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 let client_data = Arc::new(Mutex::new(ClientData {
81 state: ConnectionState::Idle,
82 queued_responses: Vec::new(),
83 memory: String::new(),
84 }));
85
86 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 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 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 if let Some(new_memory) = result.memory_updates {
141 client_data.lock().await.memory = new_memory;
142 }
143
144 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 }
176 Ok(_) => {
177 }
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 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 let mut client_data_lock = client_data.lock().await;
222
223 match client_data_lock.state {
224 ConnectionState::Idle => {
225 client_data_lock.state = ConnectionState::Processing;
227 drop(client_data_lock);
228
229 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 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 if let Some(mem) = memory_updates {
269 client_data.lock().await.memory = mem;
270 }
271
272 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 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 client_data_lock.queued_responses.push(response);
309 client_data_lock.state = ConnectionState::Accumulating;
310 }
311 ConnectionState::Accumulating => {
312 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 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}