netget/client/graphite/
mod.rs

1//! Carbon plaintext TCP emitter. The command loop owns the only socket and also handles initial actions.
2pub mod actions;
3use crate::llm::actions::client_trait::{Client, ClientActionResult};
4use crate::protocol::{ConnectContext, Event};
5use crate::server::graphite::codec;
6use crate::state::{client_handles::ClientSendOutcome, AccessLogOwner, ClientStatus};
7use crate::{console_error, console_info};
8pub use actions::GraphiteClientProtocol;
9use anyhow::{bail, Context, Result};
10use serde_json::{json, Value};
11use std::time::Duration;
12use tokio::{
13    io::{AsyncReadExt, AsyncWriteExt},
14    net::TcpStream,
15};
16const IO_TIMEOUT: Duration = Duration::from_secs(10);
17
18pub struct GraphiteClient;
19impl GraphiteClient {
20    pub async fn connect(ctx: ConnectContext) -> Result<std::net::SocketAddr> {
21        let socket = tokio::time::timeout(IO_TIMEOUT, TcpStream::connect(&ctx.remote_addr))
22            .await
23            .context("Graphite connect deadline exceeded")??;
24        let peer = socket.peer_addr()?;
25        let local_addr = socket.local_addr()?;
26        let (mut reader, mut writer) = socket.into_split();
27        let mut commands =
28            crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
29                .await;
30        ctx.state
31            .update_client_status(ctx.client_id, ClientStatus::Connected)
32            .await;
33        console_info!(
34            ctx.status_tx,
35            "Graphite emitter {} connected to {}",
36            ctx.client_id,
37            peer
38        );
39        let registrar = ctx.state.clone();
40        let client_id = ctx.client_id;
41        let handle = tokio::spawn(async move {
42            let protocol = GraphiteClientProtocol::new();
43            let connected = async {
44                let event = Event::new(
45                    &actions::GRAPHITE_CONNECTED_EVENT,
46                    json!({"remote_addr":peer.to_string(),"local_addr":local_addr.to_string()}),
47                );
48                let instruction = ctx
49                    .state
50                    .get_instruction_for_client(client_id)
51                    .await
52                    .unwrap_or_default();
53                let configured = ctx
54                    .state
55                    .get_client_event_handler_config(client_id)
56                    .await
57                    .is_some_and(|c| c.find_handler("graphite_connected").is_some());
58                if instruction.is_empty() && !configured {
59                    return Ok(crate::llm::ClientLlmResult {
60                        actions: vec![],
61                        memory_updates: None,
62                    });
63                }
64                let memory = ctx
65                    .state
66                    .get_memory_for_client(client_id)
67                    .await
68                    .unwrap_or_default();
69                crate::client::llm_budget::call_llm_for_client(
70                    &ctx.llm_client,
71                    &ctx.state,
72                    client_id.to_string(),
73                    &instruction,
74                    &memory,
75                    Some(&event),
76                    &protocol,
77                    &ctx.status_tx,
78                )
79                .await
80            };
81            tokio::pin!(connected);
82            let mut connected_done = false;
83            let mut unexpected = [0u8; 1];
84            'session: loop {
85                tokio::select! {
86                    result = reader.read(&mut unexpected) => {
87                        match result {
88                            Ok(0) => { console_info!(ctx.status_tx, "Graphite collector closed the connection"); },
89                            Ok(_) => { console_error!(ctx.status_tx, "Graphite collector sent unexpected response data; closing"); },
90                            Err(error) => { console_error!(ctx.status_tx, "Graphite collector read failed: {}", error); }
91                        }
92                        break;
93                    }
94                    result = &mut connected, if !connected_done => {
95                        connected_done = true;
96                        match result {
97                            Ok(result) => {
98                                if let Some(memory) = result.memory_updates { ctx.state.set_memory_for_client(client_id, memory).await; }
99                                for action in result.actions {
100                                    match Self::apply(&mut writer, &protocol, action).await {
101                                        Ok(ClientSendOutcome::Disconnected) => break 'session,
102                                        Ok(ClientSendOutcome::Rejected { error }) => { console_error!(ctx.status_tx, "Graphite initial action rejected: {}", error); }
103                                        Ok(_) => {},
104                                        Err(error) => { console_error!(ctx.status_tx, "Graphite send failed: {}", error); break 'session; }
105                                    }
106                                }
107                            }
108                            Err(error) => { console_error!(ctx.status_tx, "Graphite connected handler failed: {}", error); }
109                        }
110                    }
111                    command = commands.recv() => {
112                        let Some(command) = command else { break; };
113                        let action = command.action.clone();
114                        let result = Self::apply(&mut writer, &protocol, action.clone()).await;
115                        let disconnect = result.is_err() || matches!(result, Ok(ClientSendOutcome::Disconnected));
116                        let response = match &result { Ok(outcome) => serde_json::to_value(outcome).unwrap_or(Value::Null), Err(error) => json!({"error":error.to_string()}) };
117                        ctx.state.record_access_log(AccessLogOwner::Client(client_id.as_u32()), "Graphite", None, "injected_action", action, vec![response]).await;
118                        crate::client::command_support::reply(command, result);
119                        if disconnect { break; }
120                    }
121                }
122                let _ = ctx.status_tx.send("__UPDATE_UI__".into());
123            }
124            // Dropping the pending connected-handler future also cancels it when
125            // disconnect arrives while a manual/model response is outstanding.
126            ctx.state.remove_client_handle(client_id).await;
127            ctx.state
128                .update_client_status(client_id, ClientStatus::Disconnected)
129                .await;
130            let _ = ctx.status_tx.send("__UPDATE_UI__".into());
131        });
132        registrar.register_client_task(client_id, handle).await;
133        Ok(local_addr)
134    }
135    async fn apply(
136        writer: &mut tokio::net::tcp::OwnedWriteHalf,
137        protocol: &GraphiteClientProtocol,
138        action: Value,
139    ) -> Result<ClientSendOutcome> {
140        let result = match protocol.execute_action(action) {
141            Ok(result) => result,
142            Err(error) => {
143                return Ok(ClientSendOutcome::Rejected {
144                    error: error.to_string(),
145                })
146            }
147        };
148        match result {
149            ClientActionResult::Custom { name, data } if name == "send_graphite_batch" => {
150                let metrics: Vec<codec::Metric> = serde_json::from_value(data)?;
151                let bytes = match codec::encode_batch(&metrics) {
152                    Ok(bytes) => bytes,
153                    Err(error) => {
154                        return Ok(ClientSendOutcome::Rejected {
155                            error: error.to_string(),
156                        })
157                    }
158                };
159                tokio::time::timeout(IO_TIMEOUT, writer.write_all(&bytes)).await.context("Graphite write deadline exceeded; stream closed to avoid replay after partial write")??;
160                let bytes_sent = bytes.len();
161                Ok(ClientSendOutcome::Sent { bytes_sent })
162            }
163            ClientActionResult::Disconnect => Ok(ClientSendOutcome::Disconnected),
164            _ => bail!("unsupported Graphite action result"),
165        }
166    }
167}
netget/client/graphite/mod.rs

netget/client/graphite/
mod.rs
