netget/server/
ics_support.rs

1//! Shared task ownership and bounded I/O for request/reply device simulators.
2//! Wire formats and decisions remain in each protocol's session implementation.
3use crate::{
4    llm::actions::{client_trait::Client, protocol_trait::Server},
5    protocol::{ConnectContext, Event, SpawnContext},
6    server::connection::ConnectionId,
7    state::{AccessLogOwner, ClientStatus},
8};
9use anyhow::{ensure, Result};
10use async_trait::async_trait;
11use serde_json::Value;
12use std::{net::SocketAddr, sync::Arc, time::Duration};
13use tokio::{
14    io::{AsyncReadExt, AsyncWriteExt},
15    net::TcpStream,
16};
17
18pub const FIRST_BYTE_TIMEOUT: Duration = Duration::from_secs(30);
19pub const IDLE_TIMEOUT: Duration = Duration::from_secs(600);
20pub const RESPONSE_TIMEOUT: Duration = Duration::from_secs(10);
21
22pub fn parameter(
23    name: &str,
24    kind: &str,
25    description: &str,
26    required: bool,
27) -> crate::llm::actions::Parameter {
28    crate::llm::actions::Parameter {
29        name: name.into(),
30        type_hint: kind.into(),
31        description: description.into(),
32        required,
33    }
34}
35pub fn action(
36    name: &str,
37    description: &str,
38    parameters: Vec<crate::llm::actions::Parameter>,
39    example: Value,
40) -> crate::llm::actions::ActionDefinition {
41    crate::llm::actions::ActionDefinition {
42        name: name.into(),
43        description: description.into(),
44        parameters,
45        example,
46        log_template: Some(crate::protocol::LogTemplate::new().with_info(name)),
47    }
48}
49pub fn number(v: &Value, key: &str, max: u64) -> Result<u64> {
50    let n = v[key]
51        .as_u64()
52        .ok_or_else(|| anyhow::anyhow!("{key} must be an unsigned integer"))?;
53    ensure!(n <= max, "{key} exceeds {max}");
54    Ok(n)
55}
56/// A bounded framed stream read. The header determines the whole length, validated before allocation.
57pub async fn frame(
58    stream: &mut TcpStream,
59    header_len: usize,
60    max: usize,
61    length: fn(&[u8]) -> Result<usize>,
62) -> Result<Vec<u8>> {
63    let mut header = vec![0; header_len];
64    stream.read_exact(&mut header).await?;
65    let len = length(&header)?;
66    ensure!(len >= header_len && len <= max, "invalid frame length");
67    header.resize(len, 0);
68    stream.read_exact(&mut header[header_len..]).await?;
69    Ok(header)
70}
71#[async_trait]
72pub trait DeviceSession: Send + 'static {
73    async fn read(&mut self, stream: &mut TcpStream) -> Result<Vec<u8>>;
74    /// Automatic transport replies or a semantic request for the handler.
75    fn receive(&mut self, frame: &[u8]) -> Result<(Vec<u8>, Option<Value>)>;
76    fn answer(&mut self, request: &Value, action: Option<&Value>) -> Result<Vec<u8>>;
77}
78#[async_trait]
79pub trait ScannerSession: Send + 'static {
80    async fn open(&mut self, stream: &mut TcpStream) -> Result<()>;
81    async fn exchange(&mut self, stream: &mut TcpStream, action: &Value) -> Result<Value>;
82    async fn idle(&mut self, stream: &mut TcpStream) -> Result<Option<Value>> {
83        let mut b = [0];
84        stream.peek(&mut b).await?;
85        anyhow::bail!("unexpected data or peer closed")
86    }
87}
88pub async fn spawn<S: DeviceSession, F: Fn() -> S + Send + Sync + Clone + 'static>(
89    ctx: SpawnContext,
90    protocol: Arc<dyn Server>,
91    factory: F,
92    kinds: &'static [crate::protocol::EventType],
93) -> Result<SocketAddr> {
94    let listener =
95        crate::server::socket_helpers::create_reusable_tcp_listener(ctx.legacy_listen_addr())
96            .await?;
97    spawn_accept_bounded(ctx, protocol, factory, kinds, listener).await
98}
99pub async fn spawn_accept_bounded<
100    S: DeviceSession,
101    F: Fn() -> S + Send + Sync + Clone + 'static,
102>(
103    ctx: SpawnContext,
104    protocol: Arc<dyn Server>,
105    factory: F,
106    kinds: &'static [crate::protocol::EventType],
107    listener: tokio::net::TcpListener,
108) -> Result<SocketAddr> {
109    let local = listener.local_addr()?;
110    let state = ctx.state.clone();
111    let sid = ctx.server_id;
112    let task = tokio::spawn(async move {
113        let cap = crate::server::accept_bounded::ConnectionLimiter::new(
114            crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
115        );
116        loop {
117            let Ok((mut stream, peer, permit)) = crate::server::accept_bounded::accept_bounded(
118                &listener,
119                &cap,
120                b"",
121                protocol.protocol_name(),
122                Some(&ctx.status_tx),
123            )
124            .await
125            else {
126                break;
127            };
128            let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
129            let now = crate::utils::clock::Instant::now();
130            ctx.state
131                .add_connection_to_server(
132                    sid,
133                    crate::state::server::ConnectionState {
134                        id,
135                        remote_addr: peer,
136                        local_addr: local,
137                        bytes_sent: 0,
138                        bytes_received: 0,
139                        packets_sent: 0,
140                        packets_received: 0,
141                        last_activity: now,
142                        status: crate::state::server::ConnectionStatus::Active,
143                        status_changed_at: now,
144                        protocol_info: crate::state::server::ProtocolConnectionInfo::empty(),
145                    },
146                )
147                .await;
148            let mut commands =
149                crate::server::peer_support::register_peer_channel(&ctx.state, sid, id.as_u32())
150                    .await;
151            let factory = factory.clone();
152            let child = ctx.clone();
153            let p = protocol.clone();
154            ctx.state.spawn_server_task(sid, async move {
155                let _permit = permit; let mut session = factory(); let mut first = true;
156                let result: Result<()> = async {
157                    loop {
158                        let input = {
159                            let read=tokio::time::timeout(if first {FIRST_BYTE_TIMEOUT} else {IDLE_TIMEOUT},session.read(&mut stream));
160                            tokio::pin!(read);
161                            loop {
162                                tokio::select! {
163                                    v=&mut read=>break Some(v??),
164                                    Some(command)=commands.recv()=>{
165                                        if command.action["type"]=="disconnect" {crate::client::command_support::reply(command,Ok(crate::state::client_handles::ClientSendOutcome::Disconnected));break None;}
166                                        crate::client::command_support::reply(command,Ok(crate::state::client_handles::ClientSendOutcome::Rejected{error:"Only disconnect is valid outside a pending request".into()}));
167                                    }
168                                }
169                            }
170                        };
171                        let Some(input)=input else {break;};
172                        first = false;
173                        child.state.update_connection_stats(sid, id, Some(input.len() as u64), None, Some(1), None).await;
174                        let (automatic, request) = session.receive(&input)?;
175                        let mut reply = automatic;
176                        if let Some(request) = request {
177                            let event = Event::new(&kinds[0], request.clone());
178                            let result = crate::llm::action_helper::call_llm(&child.llm_client, &child.state, sid, Some(id), &event, p.as_ref()).await;
179                            let answer = result.as_ref().ok().and_then(|r| r.protocol_results.iter().find_map(|a| if let crate::llm::ActionResult::Custom {data,..}=a {Some(data)} else {None}));
180                            if answer.is_none() { crate::logging::emit::Log::new(Some(&child.status_tx)).warn(format!("{} decision=fail_closed_no_action", p.protocol_name())); }
181                            reply.extend(session.answer(&request, answer)?);
182                        }
183                        if !reply.is_empty() { stream.write_all(&reply).await?; child.state.update_connection_stats(sid, id, None, Some(reply.len() as u64), None, Some(1)).await; }
184                    } Ok(())
185                }.await;
186                if let Err(e) = result { crate::logging::emit::Log::new(Some(&child.status_tx)).debug(format!("{} peer ended: {e}", p.protocol_name())); }
187                let _ = stream.shutdown().await;
188                child.state.remove_peer_handle(sid, id.as_u32()).await;
189                child.state.update_connection_status(sid, id, crate::state::server::ConnectionStatus::Closed).await;
190            }).await;
191        }
192    });
193    state.register_server_task(sid, task).await;
194    Ok(local)
195}
196pub async fn connect<S: ScannerSession>(
197    ctx: ConnectContext,
198    protocol: Arc<dyn Client>,
199    mut session: S,
200    kinds: &'static [crate::protocol::EventType],
201) -> Result<SocketAddr> {
202    ensure!(!ctx.remote_addr.is_empty(), "remote_addr is required");
203    let mut stream =
204        tokio::time::timeout(RESPONSE_TIMEOUT, TcpStream::connect(&ctx.remote_addr)).await??;
205    tokio::time::timeout(RESPONSE_TIMEOUT, session.open(&mut stream)).await??;
206    let local = stream.local_addr()?;
207    let mut commands =
208        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
209    ctx.state
210        .update_client_status(ctx.client_id, ClientStatus::Connected)
211        .await;
212    let state = ctx.state.clone();
213    let cid = ctx.client_id;
214    state
215        .spawn_client_task(cid, async move {
216            // Handler work has its own task so an unanswered manual turn does not block injected commands.
217            let (events_tx, mut events_rx) = tokio::sync::mpsc::channel::<Event>(32);
218            let (actions_tx, mut actions_rx) = tokio::sync::mpsc::channel::<Value>(32);
219            let event_ctx = ctx.clone();
220            let p = protocol.clone();
221            let turns = ctx
222                .state
223                .spawn_client_task(cid, async move {
224                    while let Some(event) = events_rx.recv().await {
225                        let instruction = event_ctx
226                            .state
227                            .get_instruction_for_client(cid)
228                            .await
229                            .unwrap_or_default();
230                        let memory = event_ctx
231                            .state
232                            .get_memory_for_client(cid)
233                            .await
234                            .unwrap_or_default();
235                        match crate::client::llm_budget::call_llm_for_client(
236                            &event_ctx.llm_client,
237                            &event_ctx.state,
238                            cid.to_string(),
239                            &instruction,
240                            &memory,
241                            Some(&event),
242                            p.as_ref(),
243                            &event_ctx.status_tx,
244                        )
245                        .await
246                        {
247                            Ok(result) => {
248                                if let Some(memory) = result.memory_updates {
249                                    event_ctx.state.set_memory_for_client(cid, memory).await;
250                                }
251                                for a in result.actions {
252                                    if actions_tx.send(a).await.is_err() {
253                                        return;
254                                    }
255                                }
256                            }
257                            Err(e) => crate::logging::emit::Log::new(Some(&event_ctx.status_tx))
258                                .warn(format!("{} client handler: {e}", p.protocol_name())),
259                        }
260                    }
261                })
262                .await;
263            let _ = events_tx.try_send(Event::new(
264                &kinds[0],
265                serde_json::json!({"remote_addr":ctx.remote_addr}),
266            ));
267            loop {
268                let (action, command) = tokio::select! {
269                    Some(c) = commands.recv() => (c.action.clone(),Some(c)),
270                    Some(a) = actions_rx.recv() => (a,None),
271                    incoming = session.idle(&mut stream) => {match incoming {Ok(Some(response))=>{let _=events_tx.try_send(Event::new(&kinds[1],response));},Ok(None)=>{},Err(_)=>break};continue;},
272                    else => break,
273                };
274                if action["type"] == "disconnect" {
275                    if let Some(c) = command {
276                        crate::client::command_support::reply(
277                            c,
278                            Ok(crate::state::client_handles::ClientSendOutcome::Disconnected),
279                        );
280                    }
281                    break;
282                }
283                if let Err(e) = protocol.execute_action(action.clone()) {
284                    if let Some(c) = command {
285                        crate::client::command_support::reply(
286                            c,
287                            Ok(crate::state::client_handles::ClientSendOutcome::Rejected {
288                                error: e.to_string(),
289                            }),
290                        );
291                    }
292                    continue;
293                }
294                let result = match tokio::time::timeout(
295                    RESPONSE_TIMEOUT,
296                    session.exchange(&mut stream, &action),
297                )
298                .await
299                {
300                    Ok(r) => r,
301                    Err(_) => Err(anyhow::anyhow!("device response deadline exceeded")),
302                };
303                let terminal = result.is_err();
304                match result {
305                    Ok(response) => {
306                        ctx.state
307                            .record_access_log(
308                                AccessLogOwner::Client(cid.as_u32()),
309                                protocol.protocol_name(),
310                                None,
311                                &kinds[1].id,
312                                response.clone(),
313                                vec![],
314                            )
315                            .await;
316                        let _ = events_tx.try_send(Event::new(&kinds[1], response.clone()));
317                        if let Some(c) = command {
318                            crate::client::command_support::reply(
319                                c,
320                                Ok(crate::state::client_handles::ClientSendOutcome::Executed {
321                                    detail: response.to_string(),
322                                }),
323                            );
324                        }
325                    }
326                    Err(e) => {
327                        if let Some(c) = command {
328                            crate::client::command_support::reply(c, Err(e));
329                        }
330                    }
331                }
332                if terminal {
333                    break;
334                }
335            }
336            turns.abort();
337            let _ = stream.shutdown().await;
338            ctx.state.remove_client_handle(cid).await;
339            ctx.state
340                .update_client_status(cid, ClientStatus::Disconnected)
341                .await;
342        })
343        .await;
344    Ok(local)
345}
346
347#[async_trait]
348pub trait DatagramSession: Send + 'static {
349    async fn exchange(
350        &mut self,
351        socket: &mut tokio::net::UdpSocket,
352        action: &Value,
353    ) -> Result<Value>;
354    async fn idle(&mut self, socket: &mut tokio::net::UdpSocket) -> Result<Option<Value>> {
355        let mut b = [0];
356        socket.recv(&mut b).await?;
357        anyhow::bail!("unexpected datagram")
358    }
359}
360pub async fn connect_udp<S: DatagramSession>(
361    ctx: ConnectContext,
362    protocol: Arc<dyn Client>,
363    mut session: S,
364    kinds: &'static [crate::protocol::EventType],
365) -> Result<SocketAddr> {
366    ensure!(!ctx.remote_addr.is_empty(), "remote_addr is required");
367    let remote: SocketAddr = ctx.remote_addr.parse()?;
368    let mut stream = tokio::net::UdpSocket::bind(if remote.is_ipv4() {
369        "0.0.0.0:0"
370    } else {
371        "[::]:0"
372    })
373    .await?;
374    stream.connect(remote).await?;
375    let local = stream.local_addr()?;
376    let mut commands =
377        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
378    ctx.state
379        .update_client_status(ctx.client_id, ClientStatus::Connected)
380        .await;
381    let state = ctx.state.clone();
382    let cid = ctx.client_id;
383    state
384        .spawn_client_task(cid, async move {
385            // Handler work has its own task so an unanswered manual turn does not block injected commands.
386            let (events_tx, mut events_rx) = tokio::sync::mpsc::channel::<Event>(32);
387            let (actions_tx, mut actions_rx) = tokio::sync::mpsc::channel::<Value>(32);
388            let event_ctx = ctx.clone();
389            let p = protocol.clone();
390            let turns = ctx
391                .state
392                .spawn_client_task(cid, async move {
393                    while let Some(event) = events_rx.recv().await {
394                        let instruction = event_ctx
395                            .state
396                            .get_instruction_for_client(cid)
397                            .await
398                            .unwrap_or_default();
399                        let memory = event_ctx
400                            .state
401                            .get_memory_for_client(cid)
402                            .await
403                            .unwrap_or_default();
404                        match crate::client::llm_budget::call_llm_for_client(
405                            &event_ctx.llm_client,
406                            &event_ctx.state,
407                            cid.to_string(),
408                            &instruction,
409                            &memory,
410                            Some(&event),
411                            p.as_ref(),
412                            &event_ctx.status_tx,
413                        )
414                        .await
415                        {
416                            Ok(result) => {
417                                if let Some(memory) = result.memory_updates {
418                                    event_ctx.state.set_memory_for_client(cid, memory).await;
419                                }
420                                for a in result.actions {
421                                    if actions_tx.send(a).await.is_err() {
422                                        return;
423                                    }
424                                }
425                            }
426                            Err(e) => crate::logging::emit::Log::new(Some(&event_ctx.status_tx))
427                                .warn(format!("{} client handler: {e}", p.protocol_name())),
428                        }
429                    }
430                })
431                .await;
432            let _ = events_tx.try_send(Event::new(
433                &kinds[0],
434                serde_json::json!({"remote_addr":ctx.remote_addr}),
435            ));
436            loop {
437                let (action, command) = tokio::select! {
438                    Some(c) = commands.recv() => (c.action.clone(),Some(c)),
439                    Some(a) = actions_rx.recv() => (a,None),
440                    incoming = session.idle(&mut stream) => {match incoming {Ok(Some(response))=>{let _=events_tx.try_send(Event::new(&kinds[1],response));},Ok(None)=>{},Err(_)=>break};continue;},
441                    else => break,
442                };
443                if action["type"] == "disconnect" {
444                    if let Some(c) = command {
445                        crate::client::command_support::reply(
446                            c,
447                            Ok(crate::state::client_handles::ClientSendOutcome::Disconnected),
448                        );
449                    }
450                    break;
451                }
452                if let Err(e) = protocol.execute_action(action.clone()) {
453                    if let Some(c) = command {
454                        crate::client::command_support::reply(
455                            c,
456                            Ok(crate::state::client_handles::ClientSendOutcome::Rejected {
457                                error: e.to_string(),
458                            }),
459                        );
460                    }
461                    continue;
462                }
463                let result = match tokio::time::timeout(
464                    RESPONSE_TIMEOUT,
465                    session.exchange(&mut stream, &action),
466                )
467                .await
468                {
469                    Ok(r) => r,
470                    Err(_) => Err(anyhow::anyhow!("device response deadline exceeded")),
471                };
472                let terminal = result.is_err();
473                match result {
474                    Ok(response) => {
475                        ctx.state
476                            .record_access_log(
477                                AccessLogOwner::Client(cid.as_u32()),
478                                protocol.protocol_name(),
479                                None,
480                                &kinds[1].id,
481                                response.clone(),
482                                vec![],
483                            )
484                            .await;
485                        let _ = events_tx.try_send(Event::new(&kinds[1], response.clone()));
486                        if let Some(c) = command {
487                            crate::client::command_support::reply(
488                                c,
489                                Ok(crate::state::client_handles::ClientSendOutcome::Executed {
490                                    detail: response.to_string(),
491                                }),
492                            );
493                        }
494                    }
495                    Err(e) => {
496                        if let Some(c) = command {
497                            crate::client::command_support::reply(c, Err(e));
498                        }
499                    }
500                }
501                if terminal {
502                    break;
503                }
504            }
505            turns.abort();
506
507            ctx.state.remove_client_handle(cid).await;
508            ctx.state
509                .update_client_status(cid, ClientStatus::Disconnected)
510                .await;
511        })
512        .await;
513    Ok(local)
514}
netget/server/ics_support.rs

netget/server/
ics_support.rs
