netget/client/nut/
mod.rs

1pub mod actions;
2use crate::client::llm_budget::call_llm_for_client;
3use crate::llm::actions::client_trait::{Client, ClientActionResult};
4use crate::logging::emit::Log;
5use crate::protocol::{ConnectContext, Event};
6use crate::server::nut::wire::{self, Request};
7use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
8use crate::state::{AccessLogOwner, ClientStatus};
9pub use actions::NutClientProtocol;
10use anyhow::{Context, Result};
11use serde_json::{json, Value};
12use std::net::SocketAddr;
13use tokio::{
14    io::{AsyncWriteExt, BufReader},
15    net::TcpStream,
16    sync::mpsc,
17};
18
19pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
20    let stream = tokio::time::timeout(wire::IO_TIMEOUT, TcpStream::connect(&ctx.remote_addr))
21        .await
22        .context("NUT connect deadline")??;
23    let local = stream.local_addr()?;
24    ctx.state
25        .update_client_status(ctx.client_id, ClientStatus::Connected)
26        .await;
27    let external =
28        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
29    let (internal_tx, internal_rx) = mpsc::channel::<Value>(16);
30    let (event_tx, mut event_rx) = mpsc::channel::<Event>(16);
31    event_tx.try_send(Event::new(
32        &actions::CONNECTED_EVENT,
33        json!({"remote_addr":ctx.remote_addr}),
34    ))?;
35    let events_ctx = ctx.clone();
36    let dispatcher = tokio::spawn(async move {
37        let protocol = NutClientProtocol;
38        while let Some(event) = event_rx.recv().await {
39            let instruction = events_ctx
40                .state
41                .get_instruction_for_client(events_ctx.client_id)
42                .await
43                .unwrap_or_default();
44            let memory = events_ctx
45                .state
46                .get_memory_for_client(events_ctx.client_id)
47                .await
48                .unwrap_or_default();
49            match call_llm_for_client(
50                &events_ctx.llm_client,
51                &events_ctx.state,
52                events_ctx.client_id.to_string(),
53                &instruction,
54                &memory,
55                Some(&event),
56                &protocol,
57                &events_ctx.status_tx,
58            )
59            .await
60            {
61                Ok(result) => {
62                    if let Some(memory) = result.memory_updates {
63                        events_ctx
64                            .state
65                            .set_memory_for_client(events_ctx.client_id, memory)
66                            .await;
67                    }
68                    for action in result.actions {
69                        if internal_tx.send(action).await.is_err() {
70                            return;
71                        }
72                    }
73                }
74                Err(e) => {
75                    Log::new(Some(&events_ctx.status_tx)).warn(format!("NUT client handler: {e}"))
76                }
77            }
78        }
79    });
80    let dispatcher_abort = dispatcher.abort_handle();
81    ctx.state
82        .register_client_task(ctx.client_id, dispatcher)
83        .await;
84    let session_ctx = ctx.clone();
85    let task = tokio::spawn(async move {
86        let result = session(&session_ctx, stream, external, internal_rx, event_tx).await;
87        // A clean EOF/LOGOUT can follow the last response immediately. Let the
88        // registered dispatcher drain those events; its action channel is now closed
89        // and it owns no socket. Removing the client still aborts both tasks.
90        if result.is_err() {
91            dispatcher_abort.abort();
92        }
93        let status = match result {
94            Ok(()) => ClientStatus::Disconnected,
95            Err(e) => {
96                Log::new(Some(&session_ctx.status_tx)).warn(format!("NUT client ended: {e}"));
97                ClientStatus::Error(e.to_string())
98            }
99        };
100        session_ctx
101            .state
102            .update_client_status(session_ctx.client_id, status)
103            .await;
104        session_ctx
105            .state
106            .remove_client_handle(session_ctx.client_id)
107            .await;
108        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
109    });
110    ctx.state.register_client_task(ctx.client_id, task).await;
111    Ok(local)
112}
113async fn session(
114    ctx: &ConnectContext,
115    stream: TcpStream,
116    mut external: mpsc::Receiver<ClientCommand>,
117    mut internal: mpsc::Receiver<Value>,
118    events: mpsc::Sender<Event>,
119) -> Result<()> {
120    let (reader, mut writer) = tokio::io::split(stream);
121    let mut reader = BufReader::new(reader);
122    loop {
123        // An idle connection also notices EOF; NUT has no unsolicited server messages.
124        let (action, command) = tokio::select! {
125            command=external.recv()=>match command {Some(c)=>(c.action.clone(),Some(c)),None=>return Ok(())},
126            action=internal.recv()=>match action {Some(a)=>(a,None),None=>return Ok(())},
127            idle=tokio::io::AsyncBufReadExt::fill_buf(&mut reader)=>{
128                anyhow::ensure!(idle?.is_empty(),"Unsolicited NUT response"); return Ok(());
129            }
130        };
131        let result = NutClientProtocol.execute_action(action.clone());
132        let request = match result {
133            Ok(ClientActionResult::Custom { name, data }) if name == "nut_request" => {
134                Request::from_action(&data)
135            }
136            Ok(ClientActionResult::Disconnect) => {
137                let result = writer
138                    .shutdown()
139                    .await
140                    .map(|_| ClientSendOutcome::Disconnected)
141                    .map_err(Into::into);
142                if let Some(command) = command {
143                    crate::client::command_support::reply(command, result);
144                }
145                return Ok(());
146            }
147            Err(e) => Err(e),
148            _ => Err(anyhow::anyhow!("Unsupported action result")),
149        };
150        let request = match request {
151            Ok(r) => r,
152            Err(e) => {
153                if let Some(command) = command {
154                    crate::client::command_support::reply(
155                        command,
156                        Ok(ClientSendOutcome::Rejected {
157                            error: e.to_string(),
158                        }),
159                    );
160                }
161                continue;
162            }
163        };
164        let bytes = request.encode()?;
165        let write_result =
166            tokio::time::timeout(wire::IO_TIMEOUT, writer.write_all(bytes.as_bytes()))
167                .await
168                .context("NUT write deadline")
169                .and_then(|v| v.map_err(Into::into));
170        if command.is_some() {
171            // Never record credential values in the injected-action log.
172            ctx.state
173                .record_access_log(
174                    AccessLogOwner::Client(ctx.client_id.as_u32()),
175                    "NUT",
176                    None,
177                    "injected_action",
178                    request.public_data(),
179                    vec![json!({"sent":write_result.is_ok()})],
180                )
181                .await;
182        }
183        if let Some(command) = command {
184            crate::client::command_support::reply(
185                command,
186                match &write_result {
187                    Ok(()) => Ok(ClientSendOutcome::Sent {
188                        bytes_sent: bytes.len(),
189                    }),
190                    Err(e) => Err(anyhow::anyhow!(e.to_string())),
191                },
192            );
193        }
194        write_result?;
195        let response = {
196            let pending = wire::read_response(&mut reader, &request);
197            tokio::pin!(pending);
198            loop {
199                tokio::select! {
200                    biased;
201                    response=&mut pending => break response?,
202                    command=external.recv()=>{
203                        let Some(command)=command else { return Ok(()); };
204                        match NutClientProtocol.execute_action(command.action.clone()) {
205                            Ok(ClientActionResult::Disconnect)=>{
206                                let result=writer.shutdown().await.map(|_|ClientSendOutcome::Disconnected).map_err(Into::into);
207                                crate::client::command_support::reply(command,result);
208                                return Ok(());
209                            },
210                            result=>{
211                                let error=match result {
212                                    Err(e)=>e.to_string(),
213                                    _=>"NUT request already pending; retry after its response".into(),
214                                };
215                                crate::client::command_support::reply(command,Ok(ClientSendOutcome::Rejected{error}));
216                            }
217                        }
218                    }
219                }
220            }
221        };
222        events
223            .try_send(Event::new(
224                &actions::RESPONSE_EVENT,
225                json!({"request":request.public_data(),"response":response}),
226            ))
227            .context("NUT event queue full; consumer stalled")?;
228        if request.operation == "logout" {
229            writer.shutdown().await?;
230            return Ok(());
231        }
232    }
233}
netget/client/nut/mod.rs

netget/client/nut/
mod.rs
