netget/client/prometheus_remote_write/
mod.rs

1pub mod actions;
2pub mod transport;
3use crate::llm::actions::client_trait::{Client, ClientActionResult};
4use crate::protocol::{ConnectContext, Event};
5use crate::server::prometheus_remote_write::codec::{self, WriteBatch};
6use crate::state::{
7    client_handles::{ClientCommand, ClientSendOutcome},
8    AccessLogOwner, ClientStatus,
9};
10use crate::{console_error, console_info};
11pub use actions::PrometheusRemoteWriteClientProtocol;
12use anyhow::{Context, Result};
13use serde_json::{json, Value};
14use std::{collections::VecDeque, future::Future, pin::Pin};
15pub const MAX_QUEUED_EVENTS: usize = 32;
16pub const MAX_HANDLER_ACTIONS: usize = 32;
17pub const MAX_FOLLOWUP_DEPTH: usize = 8;
18type Handler = Pin<Box<dyn Future<Output = Result<crate::llm::ClientLlmResult>> + Send>>;
19type Exchange = Pin<Box<dyn Future<Output = Result<transport::WriteResponse>> + Send>>;
20struct Inflight {
21    exchange: Exchange,
22    command: Option<ClientCommand>,
23    depth: usize,
24}
25fn handler(ctx: ConnectContext, event: Event) -> Handler {
26    Box::pin(async move {
27        let instruction = ctx
28            .state
29            .get_instruction_for_client(ctx.client_id)
30            .await
31            .unwrap_or_default();
32        let configured = ctx
33            .state
34            .get_client_event_handler_config(ctx.client_id)
35            .await
36            .is_some_and(|c| c.find_handler(event.id()).is_some());
37        if instruction.is_empty() && !configured {
38            ctx.state
39                .record_access_log(
40                    AccessLogOwner::Client(ctx.client_id.as_u32()),
41                    "PrometheusRemoteWrite",
42                    None,
43                    event.id(),
44                    event.data.clone(),
45                    vec![],
46                )
47                .await;
48            return Ok(crate::llm::ClientLlmResult {
49                actions: vec![],
50                memory_updates: None,
51            });
52        }
53        let memory = ctx
54            .state
55            .get_memory_for_client(ctx.client_id)
56            .await
57            .unwrap_or_default();
58        crate::client::llm_budget::call_llm_for_client(
59            &ctx.llm_client,
60            &ctx.state,
61            ctx.client_id.to_string(),
62            &instruction,
63            &memory,
64            Some(&event),
65            &PrometheusRemoteWriteClientProtocol::new(),
66            &ctx.status_tx,
67        )
68        .await
69    })
70}
71pub struct PrometheusRemoteWriteClient;
72impl PrometheusRemoteWriteClient {
73    pub async fn connect(ctx: ConnectContext) -> Result<std::net::SocketAddr> {
74        let origin = transport::Origin::parse(&ctx.remote_addr)?;
75        let token = ctx
76            .startup_params
77            .as_ref()
78            .map(|p| p.get_optional_string("auth_token"))
79            .transpose()?
80            .flatten();
81        if let Some(token) = &token {
82            codec::validate_token(token)?;
83        }
84        let path = ctx
85            .startup_params
86            .as_ref()
87            .map(|p| p.get_optional_string("path"))
88            .transpose()?
89            .flatten()
90            .unwrap_or_else(|| codec::DEFAULT_PATH.into());
91        codec::validate_path(&path)?;
92        let retry_429 = ctx
93            .startup_params
94            .as_ref()
95            .map(|p| p.get_optional_bool("retry_429"))
96            .transpose()?
97            .flatten()
98            .unwrap_or(codec::DEFAULT_RETRY_429);
99        let config = transport::Config {
100            origin: origin.clone(),
101            token: token.clone(),
102            path,
103            retry_429,
104        };
105        let mut commands =
106            crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
107                .await;
108        ctx.state
109            .update_client_status(ctx.client_id, ClientStatus::Connected)
110            .await;
111        let registrar = ctx.state.clone();
112        let id = ctx.client_id;
113        console_info!(
114            ctx.status_tx,
115            "PrometheusRemoteWrite HTTP write session ready for {}",
116            origin.authority
117        );
118        registrar.spawn_client_task(id, async move {
119            let protocol = PrometheusRemoteWriteClientProtocol::new();
120            let mut current = Some((
121                handler(
122                    ctx.clone(),
123                    Event::new(
124                        &actions::REMOTE_WRITE_CONNECTED_EVENT,
125                        json!({"remote_addr":format!("http://{}",origin.authority)}),
126                    ),
127                ),
128                0,
129            ));
130            let mut events = VecDeque::new();
131            let mut actions = VecDeque::new();
132            let mut pending: Option<Inflight> = None;
133            'session: loop {
134                if current.is_none() {
135                    if let Some((event, depth)) = events.pop_front() {
136                        current = Some((handler(ctx.clone(), event), depth));
137                    }
138                }
139                if pending.is_none() {
140                    if let Some((action, depth)) = actions.pop_front() {
141                        match prepare(&protocol, action, depth, &config) {
142                            Ok(Prepared::Write(exchange)) => {
143                                pending = Some(Inflight {
144                                    exchange,
145                                    command: None,
146                                    depth,
147                                })
148                            }
149                            Ok(Prepared::Disconnect) => break,
150                            Err(error) => console_error!(
151                                ctx.status_tx,
152                                "PrometheusRemoteWrite handler action rejected: {}",
153                                error
154                            ),
155                        }
156                        continue;
157                    }
158                }
159                tokio::select! {
160                    result = async { pending.as_mut().unwrap().exchange.as_mut().await }, if pending.is_some() => {
161                        let p = pending.take().unwrap();
162                        match result {
163                            Ok(response) => {
164                                if let Some(command) = p.command {
165                                    finish(&ctx, command, Ok(ClientSendOutcome::Executed {
166                                        detail: format!("PrometheusRemoteWrite HTTP {}", response.status),
167                                    })).await;
168                                }
169                                if events.len() == MAX_QUEUED_EVENTS {
170                                    console_error!(ctx.status_tx, "PrometheusRemoteWrite response-event queue limit");
171                                    break;
172                                }
173                                events.push_back((
174                                    Event::new(
175                                        &actions::REMOTE_WRITE_RESPONSE_EVENT,
176                                        serde_json::to_value(response).unwrap_or(Value::Null),
177                                    ),
178                                    p.depth,
179                                ));
180                            }
181                            Err(error) => {
182                                if let Some(command) = p.command {
183                                    finish(&ctx, command, Err(anyhow::anyhow!("PrometheusRemoteWrite HTTP exchange failed"))).await;
184                                }
185                                console_error!(ctx.status_tx, "PrometheusRemoteWrite HTTP exchange failed: {}", error);
186                                break;
187                            }
188                        }
189                    }
190                    result = async { current.as_mut().unwrap().0.as_mut().await }, if current.is_some() => {
191                        let (_, depth) = current.take().unwrap();
192                        match result {
193                            Ok(result) => {
194                                if let Some(memory) = result.memory_updates {
195                                    ctx.state.set_memory_for_client(id, memory).await;
196                                }
197                                if result.actions.len() > MAX_HANDLER_ACTIONS
198                                    || actions.len() + result.actions.len() > MAX_HANDLER_ACTIONS {
199                                    console_error!(ctx.status_tx, "PrometheusRemoteWrite handler action count limit");
200                                    break;
201                                }
202                                actions.extend(result.actions.into_iter().map(|a| (a, depth + 1)));
203                            }
204                            Err(error) => console_error!(ctx.status_tx, "PrometheusRemoteWrite handler failed: {}", error),
205                        }
206                    }
207                    command = commands.recv() => {
208                        let Some(command) = command else { break; };
209                        let prepared = prepare(&protocol, command.action.clone(), 0, &config);
210                        match prepared {
211                            Ok(Prepared::Disconnect) => {
212                                finish(&ctx, command, Ok(ClientSendOutcome::Disconnected)).await;
213                                break 'session;
214                            }
215                            Ok(Prepared::Write(exchange)) if pending.is_none() => {
216                                pending = Some(Inflight { exchange, command: Some(command), depth: 0 });
217                            }
218                            Ok(Prepared::Write(_)) => finish(&ctx, command, Ok(ClientSendOutcome::Rejected {
219                                error: "One write is already in flight".into(),
220                            })).await,
221                            Err(error) => finish(&ctx, command, Ok(ClientSendOutcome::Rejected {
222                                error: error.to_string(),
223                            })).await,
224                        }
225                    }
226                }
227                let _ = ctx.status_tx.send("__UPDATE_UI__".into());
228            }
229            if let Some(p) = pending.take() {
230                if let Some(command) = p.command {
231                    finish(&ctx, command, Err(anyhow::anyhow!("PrometheusRemoteWrite write cancelled"))).await;
232                }
233            }
234            drop(current.take());
235            events.clear();
236            actions.clear();
237            ctx.state.remove_client_handle(id).await;
238            ctx.state
239                .update_client_status(id, ClientStatus::Disconnected)
240                .await;
241            let _ = ctx.status_tx.send("__UPDATE_UI__".into());
242        }).await;
243        Ok("0.0.0.0:0".parse().unwrap())
244    }
245}
246enum Prepared {
247    Disconnect,
248    Write(Exchange),
249}
250fn prepare(
251    protocol: &PrometheusRemoteWriteClientProtocol,
252    action: Value,
253    depth: usize,
254    config: &transport::Config,
255) -> Result<Prepared> {
256    match protocol.execute_action(action)? {
257        ClientActionResult::Disconnect => Ok(Prepared::Disconnect),
258        ClientActionResult::Custom { name, data } if name == "write_remote_samples" => {
259            anyhow::ensure!(depth <= MAX_FOLLOWUP_DEPTH, "write follow-up depth limit");
260            let batch: WriteBatch = serde_json::from_value(data).context("invalid typed write")?;
261            let body = codec::encode_batch(&batch)?;
262            Ok(Prepared::Write(Box::pin(transport::write(
263                config.clone(),
264                batch,
265                body,
266            ))))
267        }
268        _ => anyhow::bail!("unsupported PrometheusRemoteWrite action result"),
269    }
270}
271async fn finish(ctx: &ConnectContext, command: ClientCommand, result: Result<ClientSendOutcome>) {
272    let outcome = match &result {
273        Ok(outcome) => serde_json::to_value(outcome).unwrap_or(Value::Null),
274        Err(error) => json!({"error":error.to_string()}),
275    };
276    ctx.state
277        .record_access_log(
278            AccessLogOwner::Client(ctx.client_id.as_u32()),
279            "PrometheusRemoteWrite",
280            None,
281            "injected_action",
282            command.action.clone(),
283            vec![outcome],
284        )
285        .await;
286    crate::client::command_support::reply(command, result);
287}
netget/client/prometheus_remote_write/mod.rs

netget/client/prometheus_remote_write/
mod.rs
