netget/client/fluent_forward/
mod.rs

1//! One TCP owner; commands stay available while any event handler is parked.
2pub mod actions;
3use crate::llm::actions::client_trait::{Client, ClientActionResult};
4use crate::protocol::{ConnectContext, Event};
5use crate::server::fluent_forward::codec;
6use crate::state::{client_handles::ClientSendOutcome, AccessLogOwner, ClientStatus};
7use crate::{console_error, console_info};
8pub use actions::FluentForwardClientProtocol;
9use anyhow::{Context, Result};
10use base64::{engine::general_purpose::STANDARD, Engine};
11use serde_json::{json, Value};
12use std::{
13    collections::{HashMap, VecDeque},
14    future::Future,
15    pin::Pin,
16    time::Duration,
17};
18use tokio::{
19    io::{AsyncReadExt, AsyncWriteExt},
20    net::TcpStream,
21};
22pub const IO_TIMEOUT: Duration = Duration::from_secs(10);
23pub const MAX_PENDING_ACKS: usize = 32;
24pub const MAX_QUEUED_EVENTS: usize = 32;
25pub const MAX_FOLLOWUP_DEPTH: usize = 8;
26pub const MAX_HANDLER_ACTIONS: usize = 32;
27type Handler = Pin<Box<dyn Future<Output = Result<crate::llm::ClientLlmResult>> + Send>>;
28struct Pending {
29    tag: String,
30    count: usize,
31    deadline: tokio::time::Instant,
32    depth: usize,
33}
34pub struct FluentForwardClient;
35fn handler(ctx: ConnectContext, event: Event) -> Handler {
36    Box::pin(async move {
37        let instruction = ctx
38            .state
39            .get_instruction_for_client(ctx.client_id)
40            .await
41            .unwrap_or_default();
42        let configured = ctx
43            .state
44            .get_client_event_handler_config(ctx.client_id)
45            .await
46            .is_some_and(|c| c.find_handler(event.id()).is_some());
47        if instruction.is_empty() && !configured {
48            ctx.state
49                .record_access_log(
50                    AccessLogOwner::Client(ctx.client_id.as_u32()),
51                    "FluentForward",
52                    None,
53                    event.id(),
54                    event.data.clone(),
55                    vec![],
56                )
57                .await;
58            return Ok(crate::llm::ClientLlmResult {
59                actions: vec![],
60                memory_updates: None,
61            });
62        }
63        let memory = ctx
64            .state
65            .get_memory_for_client(ctx.client_id)
66            .await
67            .unwrap_or_default();
68        crate::client::llm_budget::call_llm_for_client(
69            &ctx.llm_client,
70            &ctx.state,
71            ctx.client_id.to_string(),
72            &instruction,
73            &memory,
74            Some(&event),
75            &FluentForwardClientProtocol::new(),
76            &ctx.status_tx,
77        )
78        .await
79    })
80}
81impl FluentForwardClient {
82    pub async fn connect(ctx: ConnectContext) -> Result<std::net::SocketAddr> {
83        let socket = tokio::time::timeout(IO_TIMEOUT, TcpStream::connect(&ctx.remote_addr))
84            .await
85            .context("Forward connect deadline exceeded")??;
86        let peer = socket.peer_addr()?;
87        let local = socket.local_addr()?;
88        let (mut reader, mut writer) = socket.into_split();
89        let mut commands =
90            crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
91                .await;
92        ctx.state
93            .update_client_status(ctx.client_id, ClientStatus::Connected)
94            .await;
95        let registrar = ctx.state.clone();
96        let client_id = ctx.client_id;
97        console_info!(ctx.status_tx, "Forward emitter connected to {}", peer);
98        let task = tokio::spawn(async move {
99            let protocol = FluentForwardClientProtocol::new();
100            let mut decoder = codec::Decoder::default();
101            let mut buffer = [0; 8192];
102            let mut pending = HashMap::<String, Pending>::new();
103            let mut events = VecDeque::new();
104            let connected = Event::new(
105                &actions::FORWARD_CONNECTED_EVENT,
106                json!({"remote_addr":peer.to_string(),"local_addr":local.to_string()}),
107            );
108            let mut current = Some((handler(ctx.clone(), connected), 0));
109            let mut expiry = tokio::time::interval(Duration::from_millis(100));
110            'session: loop {
111                if current.is_none() {
112                    if let Some((event, depth)) = events.pop_front() {
113                        current = Some((handler(ctx.clone(), event), depth));
114                    }
115                }
116                tokio::select! {
117                    read = reader.read(&mut buffer) => {
118                        let n = match read {
119                            Ok(0) => break,
120                            Ok(n) => n,
121                            Err(error) => {
122                                console_error!(ctx.status_tx, "Forward read failed: {}", error);
123                                break;
124                            }
125                        };
126                        if let Err(error) = decoder.feed(&buffer[..n]) {
127                            console_error!(ctx.status_tx, "Forward ACK framing failed: {}", error);
128                            break;
129                        }
130                        loop {
131                            let node = match decoder.next_node() {
132                                Ok(Some(node)) => node,
133                                Ok(None) => break,
134                                Err(error) => {
135                                    console_error!(ctx.status_tx, "Forward ACK invalid: {}", error);
136                                    break 'session;
137                                }
138                            };
139                            let token = match codec::parse_ack(node) {
140                                Ok(token) => token,
141                                Err(error) => {
142                                    console_error!(ctx.status_tx, "Forward ACK invalid or unsupported authentication greeting: {}", error);
143                                    break 'session;
144                                }
145                            };
146                            let Some(p) = pending.remove(&token) else {
147                                console_error!(ctx.status_tx, "Forward unmatched/repeated ACK; closing");
148                                break 'session;
149                            };
150                            if p.deadline <= tokio::time::Instant::now() {
151                                Self::record_ack_timeout(&ctx, pending.len() + 1).await;
152                                break 'session;
153                            }
154                            if events.len() == MAX_QUEUED_EVENTS {
155                                console_error!(ctx.status_tx, "Forward response-event queue limit; closing");
156                                break 'session;
157                            }
158                            events.push_back((Event::new(&actions::FORWARD_ACK_EVENT,
159                                json!({"tag":p.tag,"record_count":p.count})), p.depth));
160                        }
161                    }
162                    result = async { current.as_mut().unwrap().0.as_mut().await }, if current.is_some() => {
163                        let (_, depth) = current.take().unwrap();
164                        match result {
165                            Ok(result) => {
166                                if let Some(memory) = result.memory_updates {
167                                    ctx.state.set_memory_for_client(client_id, memory).await;
168                                }
169                                if result.actions.len() > MAX_HANDLER_ACTIONS {
170                                    console_error!(ctx.status_tx, "Forward handler action count limit");
171                                    break;
172                                }
173                                for action in result.actions {
174                                    match Self::apply(&mut writer, &protocol, &mut pending, depth + 1, action).await {
175                                        Ok(ClientSendOutcome::Disconnected) => break 'session,
176                                        Ok(ClientSendOutcome::Rejected { error }) => {
177                                            console_error!(ctx.status_tx, "Forward action rejected: {}", error);
178                                        }
179                                        Ok(_) => {}
180                                        Err(error) => {
181                                            console_error!(ctx.status_tx, "Forward send failed: {}", error);
182                                            break 'session;
183                                        }
184                                    }
185                                }
186                            }
187                            Err(error) => {
188                                console_error!(ctx.status_tx, "Forward event handler failed: {}", error);
189                            }
190                        }
191                    }
192                    command = commands.recv() => {
193                        let Some(command) = command else { break; };
194                        let action = command.action.clone();
195                        let result = Self::apply(&mut writer, &protocol, &mut pending, 0, action.clone()).await;
196                        let stop = result.is_err() || matches!(result, Ok(ClientSendOutcome::Disconnected));
197                        let response = match &result {
198                            Ok(outcome) => serde_json::to_value(outcome).unwrap_or(Value::Null),
199                            Err(error) => json!({"error":error.to_string()}),
200                        };
201                        ctx.state.record_access_log(AccessLogOwner::Client(client_id.as_u32()),
202                            "FluentForward", None, "injected_action", action, vec![response]).await;
203                        crate::client::command_support::reply(command, result);
204                        if stop { break; }
205                    }
206                    _ = expiry.tick() => {
207                        if pending.values().any(|p| p.deadline <= tokio::time::Instant::now()) {
208                            Self::record_ack_timeout(&ctx, pending.len()).await;
209                            break;
210                        }
211                    }
212                }
213                let _ = ctx.status_tx.send("__UPDATE_UI__".into());
214            }
215            ctx.state.remove_client_handle(client_id).await;
216            ctx.state
217                .update_client_status(client_id, ClientStatus::Disconnected)
218                .await;
219            let _ = ctx.status_tx.send("__UPDATE_UI__".into());
220        });
221        registrar.register_client_task(client_id, task).await;
222        Ok(local)
223    }
224    async fn record_ack_timeout(ctx: &ConnectContext, pending_count: usize) {
225        ctx.state
226            .record_access_log(
227                AccessLogOwner::Client(ctx.client_id.as_u32()),
228                "FluentForward",
229                None,
230                "forward_ack_timeout",
231                json!({"pending_count":pending_count}),
232                vec![],
233            )
234            .await;
235        console_error!(
236            ctx.status_tx,
237            "Forward ACK deadline exceeded; closing without retry"
238        );
239    }
240    async fn apply(
241        writer: &mut tokio::net::tcp::OwnedWriteHalf,
242        protocol: &FluentForwardClientProtocol,
243        pending: &mut HashMap<String, Pending>,
244        depth: usize,
245        action: Value,
246    ) -> Result<ClientSendOutcome> {
247        let result = match protocol.execute_action(action) {
248            Ok(r) => r,
249            Err(error) => {
250                return Ok(ClientSendOutcome::Rejected {
251                    error: error.to_string(),
252                })
253            }
254        };
255        match result {
256            ClientActionResult::Disconnect => Ok(ClientSendOutcome::Disconnected),
257            ClientActionResult::Custom { name, data } if name == "send_forward_batch" => {
258                let batch: codec::Batch = serde_json::from_value(data)?;
259                if batch.require_ack
260                    && (pending.len() == MAX_PENDING_ACKS || depth > MAX_FOLLOWUP_DEPTH)
261                {
262                    return Ok(ClientSendOutcome::Rejected {
263                        error: "Forward pending ACK/follow-up depth limit".into(),
264                    });
265                }
266                let chunk = batch
267                    .require_ack
268                    .then(|| STANDARD.encode(rand::random::<[u8; 16]>()));
269                let bytes = match codec::encode_batch(&batch, chunk.as_deref()) {
270                    Ok(b) => b,
271                    Err(error) => {
272                        return Ok(ClientSendOutcome::Rejected {
273                            error: error.to_string(),
274                        })
275                    }
276                };
277                tokio::time::timeout(IO_TIMEOUT,writer.write_all(&bytes)).await.context("Forward write deadline exceeded; stream closed after possible partial write")??;
278                if let Some(chunk) = chunk {
279                    pending.insert(
280                        chunk,
281                        Pending {
282                            tag: batch.tag,
283                            count: batch.entries.len(),
284                            deadline: tokio::time::Instant::now() + IO_TIMEOUT,
285                            depth,
286                        },
287                    );
288                }
289                Ok(ClientSendOutcome::Sent {
290                    bytes_sent: bytes.len(),
291                })
292            }
293            _ => anyhow::bail!("unsupported Forward action result"),
294        }
295    }
296}
netget/client/fluent_forward/mod.rs

netget/client/fluent_forward/
mod.rs
